PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: fireheart17
Tester: kingmessi
Editorialist: iceknight1093
DIFFICULTY:
easy-med
PREREQUISITES:
DFS, binary lifting
PROBLEM:
You’re given a tree. Answer Q queries on it.
Each query is as follows:
- You’re given A and B (A \neq B). A monster starts at A and moves back-and-forth along the A\leftrightarrow B path, one step at a time.
- Find the number of pairs of vertices (u, v) such that you can start at u and reach v, without ever being on the same vertex or edge at the monster at any point of time. You’re allowed to move to an adjacent node or wait, every second.
EXPLANATION:
Our first order of business is of course to answer a single query.
So, suppose we know A and B.
Let’s fix vertices u and v, and try to understand when v can be reached from this point.
Let P_{uv} denote the path from u to v, and P_{AB} the path from A to B.
First off, if the two paths don’t intersect at all we don’t have to worry about the monster and can always reach v.
So, let’s look at the case where they do intersect.
There are two possibilities: the intersection equals the entirety of P_{AB}, or it does not.
First, suppose the intersection is not the entirety of P_{AB}.
Then, it turns out it’s (almost) always possible to reach v from u.
How?
There’s a fairly simple strategy to do so.
First, move from u towards P_{AB}, till you’re just one vertex off of P_{AB}.
Now, you’ll want to move in only one direction along P_{AB}. Wait at this vertex till the monster is also moving in that direction, and is one vertex away from you on the path.
Then, simply ‘follow’ the monster, staying one step behind it always.Since the intersection isn’t the full path, you will eventually reach the point you need to leave P_{AB} before the monster reaches the end and changes direction; so you can just leave the then reach v from there.
There’s only one issue with this strategy: it assumed that we had a vertex outside the path to wait at for the monster to be in the right position and in the right direction, which we can’t always guarantee: in particular, if we start on the path itself we might have an issue with this.
We’ll come back to this case later.
That leaves us with the case when the intersection is indeed the entirety of P_{AB}.
Without loss of generality, suppose the path is u\to A\to B\to v.
If we try mimicking what was done in the previous case, we can first get as close to A as possible; then wait for the monster to return to A, and then follow one step behind it till it reaches B.
However, we now run into an issue - the monster turns around, and we have nowhere to go so we must turn around as well to avoid it. Repeating this would put us right back at A (and then beyond it) which is right where we started.
There is, however, one way to avoid backtracking the entire path.
If there’s some x\in P_{AB}\setminus \{A, B\} (i.e. a non-endpoint of the path) such that x has a neighbor not on P_{AB}, we can instead move to this neighbor when we’re at x, allow the monster to pass, and then cross it after which we’re free to go straight to v.
One way to think of this is that moving from u to x is possible from the first case (the intersection of their paths is not P_{AB}), and similarly moving from x to v is always possible; so we’re just chaining these moves together.
So, if such an x exists then u can reach v; otherwise it cannot.
Note that such a vertex exists on P_{AB} if and only if it has degree \gt 2.
This is because if its degree is \gt 2 then it will definitely have a neighbor not on the path; and if its degree is two we know both of its neighbors are on the path since it’s not an endpoint.
Thus, if P_{AB}\cap P_{uv} = P_{AB} then u can reach v only if there exists an x\in P_{AB}\setminus \{A,B\} with \text{deg}(x) \gt 2.
Note that we once again run into the same problem as the first case: if we start on the path itself, some of this analysis falls apart.
Now, suppose u lies on P_{AB} itself.
Note that if B is not a leaf, then we can move towards B and then out of the path; at which point the analysis becomes the same as above: if there’s a vertex with degree \gt 2 on the path then every v can be reached, otherwise we can’t reach exactly those vertices that are on the “other side” of P_{AB}, i.e. the set of vertices we can reach is the same as that of B.
However, what if B is a leaf?
In that case, we have two possibilities:
- There’s no vertex with degree \gt 2 on the path.
Here, we’re rather stuck, and have no way to get past the monster at all.
The only vertices we can reach are those that are closer to us than to the monster: which is everything ‘behind’ us and everything that’s at most half the distance from u to A.
Summing this across all u on the path can be done using some simple algebra, depending purely on the length of the path. - There is some vertex with degree \gt 2 on the path.
If we can reach this without getting caught by the monster then great - every v can be reached.
Otherwise, we’re functionally in the previous case, and can only reach some vertices along the path.
For the second case, we need to find the closest vertex to B along the path that has degree \gt 2, which can be done using binary lifting.
Once this distance is known, once again a bit of algebra work will give us the number of vertices that can be visited.
Now we know how to check for a single (u, v) pair, so we move to counting all possible pairs.
There are two possibilities:
- There exists x\in P_{AB} \setminus \{A, B\} such that \text{deg}(x) \gt 2.
- In this case, any (u, v) pair is valid, as long as u\neq A. The answer is just (N-1)^2.
- No such x exists.
- Here, we must start with all valid pairs and then subtract invalid ones.
- Note that a pair (u, v) is invalid if u and v are on opposite sides of A (with respect to the path).
- The number of such paths can be computed with a bit of casework, using subtree sizes.
- The cases that need to be considered are whether one of A or B is an ancestor of the other; or neither are.
For the ancestor cases, you’ll need to find which child of the higher vertex contains the lower one.
- The cases that need to be considered are whether one of A or B is an ancestor of the other; or neither are.
Finding the child of A that contains B (or vice versa) can be done in several ways; the simplest is perhaps to just binary lift from the lower node while ensuring you don’t cross the higher one.
Observe that if subtree sizes are known, then as soon as we know whether a valid x exists the actual counting can be done quickly.
So, all that remains is actually finding such an x.
That ends up being quite easy.
Let’s mark all vertices with degree \gt 2 with the value 1, and other vertices with the value 0.
Then all we’re looking for is the path sum between A and B, minus the values at A and B.
This is now a standard problem and can be solved using LCA for example:
- Let a_u denote the value assigned to u. Compute b_u as the sum of values from the root to u, which can be done using a DFS.
- Then the answer for the path (A, B) is (b_A + b_B - 2b_L + a_L) - a_A - a_B, where L = \text{lca}(A, B).
Finally, don’t forget to adjust counts for the case when B is a leaf.
TIME COMPLEXITY:
\mathcal{O}((N + Q)\log N) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define int long long
vector <ll> adj[200005];
int n, l;
int timer;
vector<int> tin, tout, level, sub;
vector<vector<int>> up, upbr, lobr;
void dfs(int v, int p, int lev)
{
tin[v] = ++timer;
up[v][0] = p;
level[v]=lev;
if(adj[v].size()<3){
upbr[v][0]=-1;
lobr[v][0]=-1;
}else{
upbr[v][0]=v;
lobr[v][0]=v;
}
for (int i = 1; i <= l; ++i){
up[v][i] = up[up[v][i-1]][i-1];
if(upbr[up[v][i-1]][i-1]!=-1){
upbr[v][i]=upbr[up[v][i-1]][i-1];
}else{
upbr[v][i]=upbr[v][i-1];
}
if(lobr[v][i-1]!=-1){
lobr[v][i]=lobr[v][i-1];
}else{
lobr[v][i]=lobr[up[v][i-1]][i-1];
}
}
sub[v]=1;
for (int u : adj[v]) {
if (u != p){
dfs(u,v,lev+1);
sub[v]+=sub[u];
}
}
tout[v] = ++timer;
}
bool is_ancestor(int u, int v)
{
return tin[u] <= tin[v] && tout[u] >= tout[v];
}
int lca(int u, int v)
{
if (is_ancestor(u, v))
return u;
if (is_ancestor(v, u))
return v;
for (int i = l; i >= 0; --i) {
if (!is_ancestor(up[u][i], v))
u = up[u][i];
}
return up[u][0];
}
ll find_branch_up(ll pos,ll dist){
ll d=dist;
ll p=pos;
ll ans=-1;
for(int i=l;i>=0;i--){
if(dist&(1<<i)){
if(upbr[p][i]!=-1){
ans=upbr[p][i];
}
p=up[p][i];
}
}
return ans;
}
ll find_branch_low(ll pos,ll dist){
ll d=dist;
ll p=pos;
ll ans=-1;
for(int i=l;i>=0;i--){
if(dist&(1<<i)){
if(lobr[p][i]!=-1){
return lobr[p][i];
}
p=up[p][i];
}
}
return ans;
}
void preprocess(int root) {
tin.resize(n);
tout.resize(n);
level.resize(n);
sub.resize(n);
timer = 0;
l = ceil(log2(n));
up.assign(n, vector<int>(l + 1));
upbr.assign(n, vector<int>(l + 1));
lobr.assign(n, vector<int>(l + 1));
dfs(root, root, 0);
}
ll solve(ll d,ll x){
ll z=(d%2==0)?((x+1)/2):(x/2);
ll ans=x*d-(x*(x+1))/2-z;
return ans/2;
}
signed main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
ll kitne_cases_hain;
kitne_cases_hain=1;
cin>>kitne_cases_hain;
while(kitne_cases_hain--){
ll q;
cin>>n>>q;
for(int i=0;i<=n;i++){
adj[i].clear();
}
ll x,y;
for(int i=1;i<n;i++){
cin>>x>>y;
x--;y--;
adj[x].push_back(y);
adj[y].push_back(x);
}
for(int i=0;i<n;i++){
if(adj[i].size()==1){
preprocess(i);
break;
}
}
ll cnt,ans,bra,tt,a,b;
while(q--){
cin>>a>>b;
a--;b--;
cnt=lca(a,b);
if(cnt!=a && cnt!=b){
if(adj[b].size()>=2){
ans=(n-1)*(n-1);
}else{
bra=find_branch_low(up[b][0],level[b]-level[cnt]);
tt=level[a]-2*level[cnt]+2*level[bra];
if(tt>level[b]){
ans=(n-1)*(n-1);
}else{
tt=level[b]-tt+1;
ans=(n-1-tt)*(n-1)+solve(level[a]+level[b]-2*level[cnt],tt);
ans+=(tt*(tt-1))/2;
}
}
}else if(cnt==b){
bra=find_branch_up(up[a][0],level[a]-level[b]-1);
if(bra==-1){
ans=(sub[a]-1)*(sub[a]+level[a]-level[b]-2);
if(adj[b].size()>=2){
ans+=(n-sub[a])*(n-sub[a]-1);
}else{
ans+=solve(level[a]-level[b],level[a]-level[b]);
ans+=((level[a]-level[b])*(level[a]-level[b]-1))/2;
}
}else{
tt=2*level[bra]-level[a];
if(adj[b].size()>=2){
ans=(n-1)*(n-1);
}else{
if(tt<level[b]){
ans=(n-1)*(n-1);
}else{
tt=tt-level[b]+1;
ans=(n-1-tt)*(n-1)+solve(level[a]-level[b],tt);
ans+=(tt*(tt-1))/2;
}
}
}
}else{
bra=find_branch_low(up[b][0],level[b]-level[a]-1);
if(bra==-1){
ans=(n-sub[b]-(level[b]-level[a]))*(n-sub[b]-1);
if(adj[b].size()>=2){
ans+=(sub[b]+(level[b]-level[a])-1)*(sub[b]+(level[b]-level[a])-2);
}else{
ans+=solve(level[b]-level[a],level[b]-level[a]);
ans+=((level[b]-level[a])*(level[b]-level[a]-1))/2;
}
}else{
tt=2*level[bra]-level[a];
if(adj[b].size()>=2){
ans=(n-1)*(n-1);
}else{
if(tt>level[b]){
ans=(n-1)*(n-1);
}else{
tt=level[b]-tt+1;
ans=(n-1-tt)*(n-1)+solve(level[b]-level[a],tt);
ans+=(tt*(tt-1))/2;
}
}
}
}
cout<<ans<<"\n";
}
}
return 0;
}
Tester's code (C++)
//Har Har Mahadev
#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp> // Common file
#include <ext/pb_ds/tree_policy.hpp>
#define ll long long
// #define int long long
#define rep(i,a,b) for(int i=a;i<b;i++)
#define rrep(i,a,b) for(int i=a;i>=b;i--)
#define repin rep(i,0,n)
#define precise(i) cout<<fixed<<setprecision(i)
#define vi vector<int>
#define si set<int>
#define mii map<int,int>
#define take(a,n) for(int j=0;j<n;j++) cin>>a[j];
#define give(a,n) for(int j=0;j<n;j++) cout<<a[j]<<' ';
#define vpii vector<pair<int,int>>
#define db double
#define be(x) x.begin(),x.end()
#define pii pair<int,int>
#define pb push_back
#define pob pop_back
#define ff first
#define ss second
#define lb lower_bound
#define ub upper_bound
#define bpc(x) __builtin_popcountll(x)
#define btz(x) __builtin_ctz(x)
using namespace std;
using namespace __gnu_pbds;
typedef tree<int, null_type, less<int>, rb_tree_tag,tree_order_statistics_node_update> ordered_set;
typedef tree<pair<int, int>, null_type,less<pair<int, int> >, rb_tree_tag,tree_order_statistics_node_update> ordered_multiset;
const long long INF=1e18;
const long long M=1e9+7;
const long long MM=998244353;
int power( int N, int M){
int power = N, sum = 1;
if(N == 0) sum = 0;
while(M > 0){if((M & 1) == 1){sum *= power;}
power = power * power;M = M >> 1;}
return sum;
}
const int N = 2e5+5;
int depth[N],sub[N],val[N],w[N],up[N][20];
ll dp[N];
bool visited[N];
vi adj[N];
void dfs(int v) {
visited[v]=true;
rep(i,1,20){
if(up[v][i-1]!=-1)up[v][i] = up[up[v][i-1]][i-1];
}
sub[v] = 1;
for(int x : adj[v]) {
if(!visited[x]) {
depth[x] = depth[up[x][0] = v]+1;
val[x] += val[v];
dfs(x);
sub[v] += sub[x];
}
}
}
int jump(int x, int d) {
rep(i,0,20){
if((d >> i) & 1)
{if(x==-1)break;
x = up[x][i];}
}
return x;
}
int LCA(int a, int b) {
if(depth[a] < depth[b]) swap(a, b);
a = jump(a, depth[a] - depth[b]);
if(a == b) return a;
rrep(i,19,0){
int aT = up[a][i], bT = up[b][i];
if(aT != bT) a = aT, b = bT;
}
return up[a][0];
}
int dist(int a,int b) {
return depth[a]+depth[b]-2*depth[LCA(a,b)];
}
int sum(int a,int b){
return val[a]+val[b]-2*val[LCA(a,b)]+(w[LCA(a,b)] ? 1 : 0);
}
int n;
ll calc(int len,int sa,int sb,int br){
ll ans = 0;
if(br == -1){
if(sb > 1){
ans = 1ll*(n-sa)*(n-sa) + 1ll*(sa-1)*(n-sb);
}
else{
ans = dp[len] + 1ll*(sa-1)*(n-1);
}
}
else{
if(sb > 1){
ans = 1ll*(n-1)*n;
}
else{
int left = max(len - (2*br-2),0);
ans = 1ll*(n-1-left)*n;
if(left){
int st = len-br;
int l1 = (left+1)/2;
int l2 = left - l1;
ans += 1ll*st*(st+1) - 1ll*(st-l1)*(st-l1+1)/2 - 1ll*(st-l2)*(st-l2+1)/2;
}
}
}
ans -= (n-1); //remove (u = v) pairs
return ans;
}
void solve()
{
cin >> n;
int q;
cin >> q;
repin{
adj[i].clear();visited[i] = 0;val[i] = 0,w[i] = 0;
rep(j,0,20)up[i][j] = -1;
}
rep(i,0,n-1){
int u,v;
cin >> u >> v;
u--;v--;
adj[u].pb(v);
adj[v].pb(u);
}
repin{
if(adj[i].size() > 2)val[i] = 1,w[i] = 1;
}
dfs(0);
while(q--){
int a,b;
cin >> a >> b;
a--;b--;
int len = dist(a,b)+1;
int lc = LCA(a,b);
int sa,sb;
if(lc == a){
sb = sub[b];
int nd = jump(b,len-2);
sa = n - sub[nd];
}
else if(lc == b){
sa = sub[a];
int nd = jump(a,len-2);
sb = n - sub[nd];
}
else{
sa = sub[a];
sb = sub[b];
}
int br = -1;//stores rightmost branch between (a,b)
if(sum(a,b)-w[a]-w[b]){
if(sum(lc,b)-w[b]){
int l = dist(a,lc);
int r = len-1;
assert(r > l);
while(r > l+1){
int m = (l+r)/2;
if(sum(b,jump(b,len-1-m))-w[b])l = m;
else r = m;
}
br = r;
}
else{
int l = 0;
int r = dist(a,lc);
assert(r > l);
while(r > l+1){
int m = (l+r)/2;
if(sum(lc,jump(a,m)))l = m;
else r = m;
}
br = r;
}
}
ll ans = calc(len,sa,sb,br);
cout << ans << '\n';
}
}
signed main(){
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
#ifdef NCR
init();
#endif
#ifdef SIEVE
sieve();
#endif
dp[2] = 1;
rep(i,3,N){
dp[i] = dp[i-1] + (i-2) + i/2;
}
int t;
cin >> t;
while(t--)
solve();
return 0;
}