PROBLEM LINK:
Practice
Div-1 Contest
Div-2 Contest
Div-3 Contest
Author: Stuart Lim
Tester: Aryan Choudhary
Editorialist: Stuart Lim
DIFFICULTY:
Medium-Hard
PREREQUISITES:
Lowest common ancestor
PROBLEM:
There is a cave with N chambers and N−1 passageways forming a tree. At 0-th second, Chef is at chamber A_j. At the start of each second, he can choose to either stay in the current chamber or move along a passageway to another chamber in 1 second. He has to visit as many distinct chambers as possible before leaving the cave at chamber B_j. There is a ghost that copies Chef’s movement with a delay of K_j, i.e. the ghost enters the cave into chamber A_j at K_j-th second and at t-th second (t \geq K_j), the ghost will be in the same place that Chef was at (t-K_j)-th second. The ghost will capture the Chef if, at some moment of time, both of them are in the same chamber or using the same passageway, including the moment when Chef is about to leave at chamber B_j. For a given start chamber A_j, end chamber B_j and ghost movement delay K_j, we have to find the maximum number of distinct chambers chef can visit. We have to answer Q such queries for the same cave.
EXPLANATION:
Let us denote the query to be a, b, k.
It’s always possible to escape if you go directly from vertex a to vertex b. Of course, in most cases, this will not give the optimal answer. We handle two edge cases separately.
Edge Case 1:
The given graph is a line. We give each vertex i another label l_i and imagine the graph as a horizontal line with a new label 1 on the left and a new label n on the right. This can be finding any vertex with degree 1, relabelling it with 1, then running the depth-first search to label the other vertices. This step has time complexity \mathcal{O}(n).
We move on to answering queries. Suppose l_a \leq l_b, i.e. start vertex a is either equal to b or is to the left of b. Then, the optimal path is always to go as far left as possible, return to a before the ghost appears, then go as far right as possible before turning back to escape at vertex b before the ghost reaches vertex b. Consider three cases:
- Vertex i is to the left of vertex a. There are l_a - 1 such vertices. We can go to vertex i and return to vertex a if and only if 2 \cdot \text{dist}(i, a) < k. So, the number of vertices that can be visited is equal to \min(l_a - 1, \lfloor \frac{k - 1}{2} \rfloor).
- Vertex i is between vertices a and b inclusive. All such vertices should be visited.
- Vertex i is to the right of vertex b. Similar to case 1.
The time complexity of answering a single query is \mathcal{O}(1).
Edge case 2:
The graph is not a line, but k \leq 2. Backtracking is not possible at all, so only the vertices lying on the path from a to b can be visited. We can preprocess the tree in \mathcal{O}(n \log(n)) to find lowest common ancestor of any pair in \mathcal{O}(\log(n) which implies that we can find the number of vertices on the path from a to b in \mathcal{O}(\log(n)) run-time for any a, b.
General case:
Define a split vertex as a vertex with the degree of at least 3. Also, define a line vertex as a vertex with a degree exactly 2. Assume that we don’t fall into edge case 1 or 2. So k > 2 and our graph is not a line. So there is at least one split vertex. Notice that you can always turn back at a split vertex for k > 2. For a split vertex s with any three adjacent vertices u, v, w, the path u \rightarrow s \rightarrow v \rightarrow s \rightarrow w \rightarrow s \rightarrow u is always possible with some time spent waiting at w. Using this we can visit all split vertices and all line vertices that between two split vertices. Let us call this the safe zone.
The other vertices form several branches hanging from split vertices. Let v be a vertex on some branch and let \text{dist}(s, v) be the distance between vertex v and the nearest split vertex s. Consider the split vertex s nearest to it. Let u, w be any two vertices adjacent to s that do not lie on the path from s to v. Then, the path we will take is u \rightarrow s \rightarrow \dots \rightarrow v \rightarrow \dots \rightarrow s \rightarrow w (see Figure 2). Regardless of u and w, you must move from s to v and back to s. You must exit s before the ghost reaches s. Hence, vertex v is only reachable if 2 \cdot \text{dist}(s, v) < K. A branch with length d excluding the split vertex would therefore contribute \min(d, \lfloor \frac{k - 1}{2} \rfloor) to the answer.
Note that there are some exceptions. If vertex a is in a branch, we initially move towards the leaf vertex, return to vertex a before the ghost appears, then go to the nearest split vertex s. This allows us to visit more vertices than if we enter the branch from the split vertex. Let l be the leaf vertex of the branch containing a. The number of vertices that can be visited in
the branch is equal to \min(\text{dist}(a, l), \lfloor \frac{k - 1}{2}\rfloor) (vertices further away from s than a), plus \text{dist}(a, s) (vertices between a and s, including a but excluding s). The case where vertex b is in a branch is similar. Note that if a and b are in the same branch, we should only consider the vertex which is further away from s.
With all of this information, we can finally discuss how to implement the solution for non-line graphs (Edge Case 2 and General Case). Root the tree at some split vertex - this guarantees that the parts of the tree closer to the
root contain the safe zone, while the further parts contain the branches. Store for each node, the depth from the root and 2^k-th parent. This takes \mathcal{O}(n \log(n)) time using depth-first search and binary lifting. These values can be used to find the lowest common ancestor in \mathcal{O}(\log(n)) and answer an Edge Case 2 query in \mathcal{O}(\log(n)) runtime.
Let us denote m = \lfloor\frac{k - 1}{2}\rfloor. As seen above, a branch of length d that doesn’t contain a and b contributes \min(d, m) to the answer. Let us define c_m as the sum of
- The number of vertices in the safe zone.
- The number of vertices in branches such that the distance from the nearest split vertex is at most m.
We can calculate c_m for all m in \mathcal{O}(n) using depth-first search. The nearest split vertex and the distance to it can be computed for all vertices using 2 depth-first searches. Notice that each branch has a unique leaf vertex. Hence we can enumerate the lengths of all branches by iterating over all leaves as the length of a branch is equal to the distance from the corresponding leaf to its nearest split vertex. And using these lengths, we can compute the contribution of 2-nd type to the sum.
While computing c_m let us also store the following information,
- l_v is the unique leaf node of the branch that v is in, or equal to 0 if v is not in a branch.
- d_v is equal to the distance between vertex v and the nearest split vertex if v is in a branch and undefined otherwise.
- s_v is equal to the length of the branch containing v if v is in a branch and undefined otherwise.
Suppose the endpoint a is in some branch. We first subtract the original contribution of this branch, which is equal to \min(s_v, \lfloor\frac{k-1}{2}\rfloor). Then we add
- The number of vertices from a (including a) to the split vertex (excluding s), which is equal to d_v.
- The number of vertices from a (excluding a) to the leaf vertex in this branch (including this leaf vertex), which is equal to \min(s_v - d_v, \lfloor\frac{k-1}{2}\rfloor).
The case where b is in a branch is also similar. Care must be taken to handle the case where both endpoints are in the same branch. Each query in the General Case can be handled in \mathcal{O}(1).
Overall time complexity: \mathcal{O}(n \log(n) + q \log(n))
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
using namespace std;
int N, Q, source;
vector <int> adj[300005];
bool is_line;
// line solution only
int L[300005];
// non-line solution only
int root_depth[300005], parent[300005][20];
int safe_zone_size;
int in_branch_id[300005], in_branch_depth[300005];
int cnt_branch_depth[300005], pref_cnt_branch_depth[300005];
int branch_length[300005];
int base[300005]; // TAKES IN floor((k-1)/2) AS INDEX, NOT k
// = safe_zone_size + vertices in branches with depth up to i
void dfs_nonline_from_source(int c, int p, int d)
{
root_depth[c] = d;
parent[c][0] = p;
for (int b = 1; b < 20; b++)
parent[c][b] = parent[ parent[c][b - 1] ][b - 1];
for (int s : adj[c])
{
if (s == p) continue;
dfs_nonline_from_source(s, c, d + 1);
}
}
int lca(int u, int v)
{
if (root_depth[u] < root_depth[v]) swap(u, v);
int depth_diff = root_depth[u] - root_depth[v];
for (int b = 19; b >= 0; b--)
{
if (depth_diff & (1 << b)) u = parent[u][b];
}
if (u == v) return u;
for (int b = 19; b >= 0; b--)
{
if (parent[u][b] != parent[v][b])
{
u = parent[u][b];
v = parent[v][b];
}
}
return parent[u][0];
}
int count_vertices_on_path(int u, int v)
{
return (root_depth[u] + root_depth[v] - 2 * root_depth[lca(u, v)] + 1);
}
int main()
{
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
cin >> N >> Q;
for (int i = 1; i <= N - 1; i++)
{
int u, v;
cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
// check degree of each vertex to determine if graph is a line
is_line = 1;
for (int c = 1; c <= N; c++)
{
if (adj[c].size() > 2) is_line = 0;
}
// edge case 1 (line solution)
if (is_line)
{
// source vertex is an endpoint of the line
for (int c = 1; c <= N; c++)
{
if (adj[c].size() <= 1)
{
source = c;
break;
}
}
// fill array L
memset(L, 0, sizeof L);
L[source] = 1;
for (int c = source, i = 2; i <= N; i++)
{
if (L[adj[c][0]] == 0) c = adj[c][0];
else c = adj[c][1];
L[c] = i;
}
// answer queries
for (int q = 1; q <= Q; q++)
{
int a, b, k;
cin >> a >> b >> k;
if (L[a] > L[b]) swap(a, b);
int ans = min(L[a] - 1, (k - 1) / 2)
+ (L[b] - L[a] + 1)
+ min(N - L[b], (k - 1) / 2);
cout << ans << "\n";
}
}
// general case and edge case 2 (non-line solution)
else
{
// source vertex is any 'split' vertex
for (int c = 1; c <= N; c++)
{
if (adj[c].size() >= 3)
{
source = c;
break;
}
}
// fill depth and parent arrays
dfs_nonline_from_source(source, source, 0);
// fill in_branch_id, in_branch_depth, cnt_branch_depth
// and branch_length arrays
// calculate safe_zone_size
// start from leaf vertex and keep moving to immediate parent
// until a split vertex is found
memset(in_branch_id, 0, sizeof in_branch_id);
memset(in_branch_depth, 0, sizeof in_branch_depth);
memset(cnt_branch_depth, 0, sizeof cnt_branch_depth);
memset(branch_length, 0, sizeof branch_length);
safe_zone_size = N;
for (int c = 1; c <= N; c++)
{
if (adj[c].size() == 1)
{
int len = 0;
for (int u = c; adj[u].size() < 3; u = parent[u][0])
len++;
branch_length[c] = len;
safe_zone_size -= len;
for (int u = c, d = len; d > 0; u = parent[u][0], d--)
{
in_branch_id[u] = c;
in_branch_depth[u] = d;
cnt_branch_depth[d]++;
}
}
}
// fill in pref_cnt_branch_depth and base answer arrays
pref_cnt_branch_depth[0] = 0;
base[0] = 0;
for (int i = 1; i <= N; i++)
{
pref_cnt_branch_depth[i] = pref_cnt_branch_depth[i - 1] + cnt_branch_depth[i];
base[i] = safe_zone_size + pref_cnt_branch_depth[i];
}
// answer queries
for (int q = 1; q <= Q; q++)
{
int a, b, k, ans;
cin >> a >> b >> k;
if (k <= 2) ans = count_vertices_on_path(a, b);
else
{
ans = base[min( N, (k - 1) / 2 )];
if (in_branch_id[a] != 0 && (in_branch_id[a] != in_branch_id[b] || in_branch_depth[a] >= in_branch_depth[b]))
{
// update with contribution of a
int branch_idx = in_branch_id[a];
int base_contribution = min(branch_length[branch_idx], (k - 1) / 2);
int dist_to_split = in_branch_depth[a];
int dist_to_leaf = branch_length[branch_idx] - in_branch_depth[a];
int final_contribution = dist_to_split + min(dist_to_leaf, (k - 1) / 2);
ans = ans - base_contribution + final_contribution;
}
if (in_branch_id[b] != 0 && (in_branch_id[a] != in_branch_id[b] || in_branch_depth[a] < in_branch_depth[b]))
{
// update with contribution of b
int branch_idx = in_branch_id[b];
int base_contribution = min(branch_length[branch_idx], (k - 1) / 2);
int dist_to_split = in_branch_depth[b];
int dist_to_leaf = branch_length[branch_idx] - in_branch_depth[b];
int final_contribution = dist_to_split + min(dist_to_leaf, (k - 1) / 2);
ans = ans - base_contribution + final_contribution;
}
}
cout << ans << "\n";
}
}
return 0;
}
Tester's Solution
/* in the name of Anton */
/*
Compete against Yourself.
Author - Aryan (@aryanc403)
Atcoder library - https://atcoder.github.io/ac-library/production/document_en/
*/
#ifdef ARYANC403
#include <header.h>
#else
#pragma GCC optimize ("Ofast")
#pragma GCC target ("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx")
//#pragma GCC optimize ("-ffloat-store")
#include<bits/stdc++.h>
#define dbg(args...) 42;
#endif
// y_combinator from @neal template https://codeforces.com/contest/1553/submission/123849801
// http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2016/p0200r0.html
template<class Fun> class y_combinator_result {
Fun fun_;
public:
template<class T> explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}
template<class ...Args> decltype(auto) operator()(Args &&...args) { return fun_(std::ref(*this), std::forward<Args>(args)...); }
};
template<class Fun> decltype(auto) y_combinator(Fun &&fun) { return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun)); }
using namespace std;
#define fo(i,n) for(i=0;i<(n);++i)
#define repA(i,j,n) for(i=(j);i<=(n);++i)
#define repD(i,j,n) for(i=(j);i>=(n);--i)
#define all(x) begin(x), end(x)
#define sz(x) ((lli)(x).size())
#define pb push_back
#define mp make_pair
#define X first
#define Y second
#define endl "\n"
typedef long long int lli;
typedef long double mytype;
typedef pair<lli,lli> ii;
typedef vector<ii> vii;
typedef vector<lli> vi;
const auto start_time = std::chrono::high_resolution_clock::now();
void aryanc403()
{
#ifdef ARYANC403
auto end_time = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> diff = end_time-start_time;
cerr<<"Time Taken : "<<diff.count()<<"\n";
#endif
}
long long readInt(long long l, long long r, char endd) {
long long x=0;
int cnt=0;
int fi=-1;
bool is_neg=false;
while(true) {
char g=getchar();
if(g=='-') {
assert(fi==-1);
is_neg=true;
continue;
}
if('0'<=g&&g<='9') {
x*=10;
x+=g-'0';
if(cnt==0) {
fi=g-'0';
}
cnt++;
assert(fi!=0 || cnt==1);
assert(fi!=0 || is_neg==false);
assert(!(cnt>19 || ( cnt==19 && fi>1) ));
} else if(g==endd) {
if(is_neg) {
x=-x;
}
assert(l<=x&&x<=r);
return x;
} else {
assert(false);
}
}
}
string readString(int l, int r, char endd) {
string ret="";
int cnt=0;
while(true) {
char g=getchar();
assert(g!=-1);
if(g==endd) {
break;
}
cnt++;
ret+=g;
}
assert(l<=cnt&&cnt<=r);
return ret;
}
long long readIntSp(long long l, long long r) {
return readInt(l,r,' ');
}
long long readIntLn(long long l, long long r) {
return readInt(l,r,'\n');
}
string readStringLn(int l, int r) {
return readString(l,r,'\n');
}
string readStringSp(int l, int r) {
return readString(l,r,' ');
}
void readEOF(){
assert(getchar()==EOF);
}
vi readVectorInt(int n,lli l,lli r){
vi a(n);
for(int i=0;i<n-1;++i)
a[i]=readIntSp(l,r);
a[n-1]=readIntLn(l,r);
return a;
}
#include <algorithm>
#include <cassert>
#include <vector>
namespace atcoder {
struct dsu {
public:
dsu() : _n(0) {}
explicit dsu(int n) : _n(n), parent_or_size(n, -1) {}
int merge(int a, int b) {
assert(0 <= a && a < _n);
assert(0 <= b && b < _n);
int x = leader(a), y = leader(b);
if (x == y) return x;
if (-parent_or_size[x] < -parent_or_size[y]) std::swap(x, y);
parent_or_size[x] += parent_or_size[y];
parent_or_size[y] = x;
return x;
}
bool same(int a, int b) {
assert(0 <= a && a < _n);
assert(0 <= b && b < _n);
return leader(a) == leader(b);
}
int leader(int a) {
assert(0 <= a && a < _n);
if (parent_or_size[a] < 0) return a;
return parent_or_size[a] = leader(parent_or_size[a]);
}
int size(int a) {
assert(0 <= a && a < _n);
return -parent_or_size[leader(a)];
}
std::vector<std::vector<int>> groups() {
std::vector<int> leader_buf(_n), group_size(_n);
for (int i = 0; i < _n; i++) {
leader_buf[i] = leader(i);
group_size[leader_buf[i]]++;
}
std::vector<std::vector<int>> result(_n);
for (int i = 0; i < _n; i++) {
result[i].reserve(group_size[i]);
}
for (int i = 0; i < _n; i++) {
result[leader_buf[i]].push_back(i);
}
result.erase(
std::remove_if(result.begin(), result.end(),
[&](const std::vector<int>& v) { return v.empty(); }),
result.end());
return result;
}
private:
int _n;
std::vector<int> parent_or_size;
};
} // namespace atcoder
const lli INF = 0xFFFFFFFFFFFFFFFL;
lli seed;
mt19937 rng(seed=chrono::steady_clock::now().time_since_epoch().count());
inline lli rnd(lli l=0,lli r=INF)
{return uniform_int_distribution<lli>(l,r)(rng);}
class CMP
{public:
bool operator()(ii a , ii b) //For min priority_queue .
{ return ! ( a.X < b.X || ( a.X==b.X && a.Y <= b.Y )); }};
void add( map<lli,lli> &m, lli x,lli cnt=1)
{
auto jt=m.find(x);
if(jt==m.end()) m.insert({x,cnt});
else jt->Y+=cnt;
}
void del( map<lli,lli> &m, lli x,lli cnt=1)
{
auto jt=m.find(x);
if(jt->Y<=cnt) m.erase(jt);
else jt->Y-=cnt;
}
bool cmp(const ii &a,const ii &b)
{
return a.X<b.X||(a.X==b.X&&a.Y<b.Y);
}
#define rep(i, a, b) for(int i = a; i < (b); ++i)
#define trav(a, x) for(auto& a : x)
typedef long long ll;//lli
typedef pair<int, int> pii;
typedef vector<pii> vpi;
typedef vector<vpi> graph;
graph readTree(const int n){
graph e(n);
atcoder::dsu d(n);
for(lli i=1;i<n;++i){
const lli u=readIntSp(1,n)-1;
const lli v=readIntLn(1,n)-1;
e[u].pb({v,1});
e[v].pb({u,1});
d.merge(u,v);
}
assert(d.size(0)==n);
return e;
}
template<class T>
struct RMQ {
vector<vector<T>> jmp;
RMQ(const vector<T>& V) {
int N = sz(V), on = 1, depth = 1;
while (on < N) on *= 2, depth++;
jmp.assign(depth, V);
rep(i,0,depth-1) rep(j,0,N)
jmp[i+1][j] = min(jmp[i][j],
jmp[i][min(N - 1, j + (1 << i))]);
}
T query(int a, int b) {
assert(a < b); // or return inf if a == b
int dep = 31 - __builtin_clz(b - a);
return min(jmp[dep][a], jmp[dep][b - (1 << dep)]);
}
};
struct LCA {
vi time;
vector<ll> dist;
RMQ<pii> rmq;
LCA(graph& C) : time(sz(C), -99), dist(sz(C)), rmq(dfs(C)) {}
vpi dfs(graph& C) {
vector<tuple<int, int, int, ll>> q(1);
vpi ret;
int T = 0, v, p, d; ll di;
while (!q.empty()) {
tie(v, p, d, di) = q.back();
q.pop_back();
if (d) ret.emplace_back(d, p);
time[v] = T++;
dist[v] = di;
trav(e, C[v]) if (e.first != p)
q.emplace_back(e.first, v, d+1, di + e.second);
}
return ret;
}
int query(int a, int b) {
if (a == b) return a;
a = time[a], b = time[b];
return rmq.query(min(a, b), max(a, b)).second;
}
ll distance(int a, int b) {
int lca = query(a, b);
return dist[a] + dist[b] - 2 * dist[lca];
}
};
int main(void) {
ios_base::sync_with_stdio(false);cin.tie(NULL);
// freopen("txt.in", "r", stdin);
// freopen("txt.out", "w", stdout);
// cout<<std::fixed<<std::setprecision(35);
// readInput
const int n=readIntSp(2,2e5),q=readIntLn(1,2e5);
auto e=readTree(n);
vector<vi> queries(q,vi(3));
for(auto &v:queries){
v[0]=readIntSp(1,n)-1;
v[1]=readIntSp(1,n)-1;
v[2]=readIntLn(1,1e9);
}
readEOF();
// solveProblem
const auto isSubtask1 = [&](){
for(const auto &v:e)
if(sz(v)>=3)
return false;
return true;
};
const auto isSubtask2 = [&](){
for(const auto &v:queries)
if(v[2]!=1)
return false;
return true;
};
const auto isSubtask3 = [&](){
return n<=5000&&q<=5000;
};
const auto solveSubtask1 = [&](){
dbg("subtask 1");
const int p1=[&](){
for(int i=0;i<n;++i){
if(sz(e[i])<2)
return i;
}
assert(false);
}();
dbg(p1);
const vi position=[&]()->vi{
if(n==1)
return {0};
int u=e[p1][0].X,par=p1;
vi pos(n,-1);
pos[par]=0;
pos[u]=1;
lli i=2;
while(sz(e[u])!=1){
const int v=e[u][0].X^e[u][1].X^par;
par=u;
u=v;
pos[u]=i++;
}
return pos;
}();
dbg(position);
for(const auto &z:queries){
const int uj=position[z[0]],vj=position[z[1]];
const int u=min(uj,vj),v=max(uj,vj),k=z[2];
const int m=(k-1)/2;
cout<<v-u+1+min(u,m)+min(n-v-1,m)<<endl;
}
};
if(isSubtask1()){
solveSubtask1();
return 0;
}
// assert(isSubtask2());
// assert(isSubtask3());
LCA lca(e);
const int root = [&](){
for(int i=0;i<n;++i){
if(sz(e[i])>2)
return i;
}
assert(false);
}();
dbg(root);
// vi par(n,-1);
vi stk;stk.reserve(n);
vii hangLeaf(n,mp(-1,-1));
lli nonBranch = n;
vi prefSum(n+2,0);
y_combinator([&](const auto &self,lli u,lli p,lli h)->void{
// par[u]=p;
if(sz(e[u])>2)
h=0;
stk.pb(u);
if(sz(e[u])==1){
const int f=sz(stk)-h-1;
dbg(h);
nonBranch-=h;
prefSum[1]++;
prefSum[h+1]--;
for(int i=1;i<=h;++i)
hangLeaf[stk[f+i]]={i,u};
}
for(auto x:e[u]){
if(x.X==p)
continue;
self(x.X,u,h+1);
}
stk.pop_back();
})(root,-1,0);
dbg(nonBranch);
lli sum=0;
for(int i=1;i<=n+1;++i){
sum+=prefSum[i];
prefSum[i]=prefSum[i-1]+sum;
}
dbg(prefSum);
const auto contri=[&](lli u,lli m)->lli{
const int leaf=hangLeaf[u].Y;
if(leaf==-1)
return 0;
const int cur = min(hangLeaf[u].X+m,hangLeaf[leaf].X)-min(hangLeaf[leaf].X,m);
dbg(u,cur);
return cur;
};
for(const auto &z:queries){
const int u=z[0],v=z[1],k=z[2];
if(k<=2){
cout<<lca.distance(u,v)+1<<endl;
continue;
}
const int m=min((k-1)/2,n+1);
lli ans=nonBranch+prefSum[m];
const int leafU=hangLeaf[u].Y,leafV=hangLeaf[v].Y;
if(leafU==leafV){
if(hangLeaf[u].X>hangLeaf[v].X)
ans+=contri(u,m);
else
ans+=contri(v,m);
} else {
ans+=contri(u,m);
ans+=contri(v,m);
}
cout<<ans<<endl;
}
aryanc403();
return 0;
}