CF - Sloth Naptime

Announcement With Official Editorial

Explanation

Time Complexity: $\mathcal{O}(N \log N)$

Basically the problem statement asks if given a path from a start node to an end node on a tree, how far the sloth will travel if it has a set amount of energy, and each edge on the path has an energy cost.

We know that the sloth will always try to move upwards from the start node $\texttt{st}$ to the least common ancestor of the start and end nodes $\texttt{lca}$, and then move downwards from the least common ancestor towards the end node $\texttt{end}$.

Let $\texttt{dist}(i, j)$ be the distance from node $i$ to node $j$.

So if the sloth has energy amount $e$, there are 3 separate cases:

1. $e \geq \texttt{dist}(dist(\texttt{st}, \texttt{end})$. In this case the result will be the sloth reaching the end.

2. $e < \texttt{dist(\texttt{st}, \texttt{end})}$ and $e < \texttt{dist(\texttt{st}, \texttt{lca})}$ in which the result will be the sloth reaching the $e$th parent of $\texttt{st}$.

3. $e < \texttt{dist(\texttt{st}, \texttt{end})}$ and $e \geq \texttt{dist(\texttt{st}, \texttt{lca})}$ in which the result will be the sloth reaching the $(e - \texttt{dist(\texttt{st}, \texttt{lca}))}$th parent of $\texttt{end}$.

So first, we can can run DFS once to find the depth of every node.

Then we can create a matrix $\texttt{anc}$ to store the ancestors of each node, where $\texttt{anc[i][j]}$ is the $2^j$th ancestor of node $i$.

After that, we can use binary lifting to answer each query by finding the least common ancestor of the start and end nodes, and then finding the right node that the sloth will reach.

Implementation

Copy
import java.io.*;
import java.util.*;
public class SlothNaptime {
	public static final int MAXN = (int)3e5 + 5;
	public static final int LOGN = (int)(Math.log(MAXN) / Math.log(2)) + 1;
	public static ArrayList<Integer>[] adj = new ArrayList[MAXN];
	// anc[i][j] is the 2^j-th parent of i.
	public static int[][] anc = new int[MAXN][LOGN];
	public static int[] depth = new int[MAXN];

Copy
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 3e5 + 5;
const int LOGN = log2(MAXN) + 1;

int depth[MAXN];
// anc[i][j] is the 2^j-th parent of i.
int anc[MAXN][LOGN];
vector<int> adj[MAXN];