AT_abc267_f [ABC267F] Exactly K Steps

大家好，我是毒瘤，喜欢用玄学算法过题。

发现题解区没有这个做法，于是来发一篇。

思路

不难发现如果一个点对 $(u,v)$ 的距离为 $d$，那么在这棵树以 $u$ 为根时，$v$ 的深度为 $d$。于是考虑换根 DP。

首先思考如何计算答案。显然我们可以将查询离线下来，然后当换根到以 $u$ 为根时，将有关 $u$ 的所有查询全部解决即可。

其次，发现当一个 $v$ 转化成根时，它会在其父节点 $u$ 为根时的形态的基础上，所在 $v$ 子树中的节点深度减 $1$，其余节点加 $1$。

然而，我们需要求的是深度为 $d$ 的节点，于是单单想大多数的换根 DP 维护深度的极值是不行的，所以需要更新出所有的深度。

于是这个东西需要使用数据结构维护。明确这个数据结构所需的功能：

1. 区间加减。
2. 区间查询权值为 $k$ 的编号。

于是你就想到了这道题。即用分块实现，一个 vector 维护块内的元素，每一次每一次修改都需要一次排序，查询都需要一次二分，单次操作是 $\Theta(\sqrt n \log \sqrt n)$ 的。配合上巨大常数，你得到了 TLE 15 的代码。

然后你可以发现此题是维护深度，所以值域很小，所以可以开桶维护。你在查询的时候每次都先看一下在这块中有没有出现 $k$，如果出现了就直接暴力找；否则就继续向前面的块判断。这样实现即可单次操作做到 $\Theta(\sqrt n)$。

于是你就得到了理论时间复杂度为 $\Theta(q \sqrt n)$ 的代码，但是常数有点小大，但是发现每一次的桶都不需要清空完，只需将原本的 $val$ 删掉，然后加上新的 $val$ 即可。

Code

#include <bits/stdc++.h>  
#define fst first  
#define snd second  
#define re register  
  
using namespace std;  
  
typedef pair<int,int> pii;  
const int N = 2e5 + 10,M = 4e5 + 10,K = 1010;  
int n,q;  
int idx,h[N],ne[M],e[M],ans[N];  
int num,id[N],pid[N],sz[N],d[N],val[N];  
vector<pii> Q[N];  
  
inline int read(){  
    int r = 0,w = 1;  
    char c = getchar();  
    while (c < '0' || c > '9'){  
        if (c == '-') w = -1;  
        c = getchar();  
    }  
    while (c >= '0' && c <= '9'){  
        r = (r << 3) + (r << 1) + (c ^ 48);  
        c = getchar();  
    }  
    return r * w;  
}  
  
inline void add(int a,int b){  
    ne[idx] = h[a];  
    e[idx] = b;  
    h[a] = idx++;  
}  
  
struct block{  
    int len,num;  
    int pos[N],L[K],R[K],tag[K],vis[K][N];  
  
    inline void init(){  
        len = num = sqrt(n);  
        for (re int i = 1;i <= n;i++) val[id[i]] = d[i];  
        for (re int i = 1;i <= num;i++){  
            L[i] = (i - 1) * len + 1;  
            R[i] = i * len;  
        }  
        if (R[num] != n){  
            num++;  
            L[num] = (num - 1) * len + 1;  
            R[num] = n;  
        }  
        for (re int i = 1;i <= num;i++){  
            for (re int j = L[i];j <= R[i];j++){  
                pos[j] = i;  
                vis[i][val[j]]++;  
            }  
        }  
    }  
  
    inline void modify(int l,int r,int k){  
        int p = pos[l],q = pos[r];  
        if (p == q){  
            for (re int i = l;i <= r;i++){  
                vis[p][val[i]]--;  
                val[i] += k;  
                vis[p][val[i]]++;  
            }  
        }  
        else{  
            for (re int i = l;i <= R[p];i++){  
                vis[p][val[i]]--;  
                val[i] += k;  
                vis[p][val[i]]++;  
            }  
            for (re int i = p + 1;i < q;i++) tag[i] += k;  
            for (re int i = L[q];i <= r;i++){  
                vis[q][val[i]]--;  
                val[i] += k;  
                vis[q][val[i]]++;  
            }  
        }  
    }  
  
    inline int query(int l,int r,int k){  
        int p = pos[l],q = pos[r];  
        if (p == q){  
            for (re int i = l;i <= r;i++){  
                if (val[i] + tag[p] == k) return pid[i];  
            }  
        }  
        else{  
            for (re int i = l;i <= R[p];i++){  
                if (val[i] + tag[p] == k) return pid[i];  
            }  
            for (re int i = p + 1;i < q;i++){  
                if (vis[i][k - tag[i]]){  
                    for (re int j = L[i];j <= R[i];j++){  
                        if (val[j] + tag[i] == k) return pid[j];  
                    }  
                }  
            }  
            for (re int i = L[q];i <= r;i++){  
                if (val[i] + tag[q] == k) return pid[i];  
            }  
        }  
        return -1;  
    }  
}bl;  
  
inline void calc(int u){  
    for (auto x:Q[u]) ans[x.snd] = bl.query(1,n,x.fst + 1);  
}  
  
inline void dfs1(int u,int fa){  
    num++;  
    sz[u] = 1;  
    id[u] = num;  
    pid[num] = u;  
    d[u] = d[fa] + 1;  
    for (re int i = h[u];~i;i = ne[i]){  
        int j = e[i];  
        if (j == fa) continue;  
        dfs1(j,u);  
        sz[u] += sz[j];  
    }  
}  
  
inline void dfs2(int u,int fa){  
    for (re int i = h[u];~i;i = ne[i]){  
        int j = e[i];  
        if (j == fa) continue;  
        bl.modify(1,n,1);  
        bl.modify(id[j],id[j] + sz[j] - 1,-2);  
        calc(j);  
        dfs2(j,u);  
        bl.modify(1,n,-1);  
        bl.modify(id[j],id[j] + sz[j] - 1,2);  
    }  
}  
  
int main(){  
    memset(h,-1,sizeof(h));  
    n = read();  
    for (re int i = 1;i < n;i++){  
        int a,b;  
        a = read();  
        b = read();  
        add(a,b);  
        add(b,a);  
    }  
    q = read();  
    for (re int i = 1;i <= q;i++){  
        int a,b;  
        a = read();  
        b = read();  
        Q[a].push_back({b,i});  
    }  
    for (re int i = 1;i <= n;i++) sort(Q[i].begin(),Q[i].end());  
    dfs1(1,0);  
    bl.init();  
    calc(1);  
    dfs2(1,0);  
    for (re int i = 1;i <= q;i++) printf("%d\n",ans[i]);  
    return 0;  
}