Codeforces 1585E Frequency Queries

题目大意

给定一棵 \(n(n\leq 10^6)\) 个点的树，结点 \(u\) 的点权为 \(a_u(1\leq a_u\leq n)\)，\(q(q\leq 10^6)\) 次询问，每次询问 \(u,l,k\)，表示从 \(u\) 到根的树链上所有点权拿出来，把点权按出现次数从小到大排序再去重，去掉出现次数少于 \(l\) 的点权，剩下的数组里的第 \(k\) 个点权是多少。答案可能有多解，输出任意一个，若无解，输出 \(-1\)。

题解

首先询问可以离线，所以可以每遍历到一个点就处理关于这个点的所有询问。

设 \(cnt[x]\) 表示点权 \(x\) 出现的数目。在DFS的过程中很容易可以维护当前点 \(u\) 到根的路径上所有的点权的 \(cnt[x]\)。

设 \(p\) 为 \(1\sim n\) 的一个排列，表示点权 \(1\sim n\) 按在 \(u\) 到根的路径上出现的次数从小到大排序的结果，设 \(p^{-1}\) 表示 \(p\) 的逆排列，即 \(p^{-1}[p[i]]=i\)，设 \(lb_x\) 表示第一个出现次数大于等于 \(x\) 的点权在 \(p\) 中的下标。我们只要维护 \(cnt,p,p^{-1},lb\) 数组，最后询问 \(u,l,k\) 的答案即为 \(p[lb[l]+k-1]\)。

当DFS进入一个点 \(u\) 时，我们令 \(x:=a[u]\)。此时 \(p\) 一定是被分为若干段，每一段内的点权出现的次数相等，段与段之间点权出现的次数相差1。那么我们需要考虑 \(++cnt[x]\) 时对 \(p,p^{-1},lb\) 的影响。我们可以把 \(p\) 中 \(x\) 的位置和 \(x\) 所在段的最后一个数的位置互换，然后让当前段的最后一个位置变成下一段的第一个位置即可，即：

int x = a[u];
int p1 = pt[x], p2 = lb[cnt[x] + 1] - 1;
swap(p[p1], p[p2]);
int pt1 = x, pt2 = p[pt[x]];
swap(pt[pt1], pt[pt2]);
--lb[cnt[x] + 1]; ++cnt[x];

当退出 \(u\) 结点时需要恢复影响：

--cnt[x]; ++lb[cnt[x] + 1];
swap(pt[pt1], pt[pt2]);
swap(p[p1], p[p2]);

时间复杂度 \(O(n+q)\)。

感觉这个单点 \(O(1)\) 修改挺巧妙的。

Code

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

template<typename elemType>
inline void Read(elemType& T) {
    elemType X = 0, w = 0; char ch = 0;
    while (!isdigit(ch)) { w |= ch == '-';ch = getchar(); }
    while (isdigit(ch)) X = (X << 3) + (X << 1) + (ch ^ 48), ch = getchar();
    T = (w ? -X : X);
}

const int maxn = 1000010;

vector<int> G[maxn];
struct Node { int id, l, k; };
vector<Node> ask[maxn];
int a[maxn], ans[maxn], cnt[maxn], p[maxn], pt[maxn], lb[maxn];
int T, n, q;

void DFS(int u, int fa) {
    int x = a[u];
    int p1 = pt[x], p2 = lb[cnt[x] + 1] - 1;
    swap(p[p1], p[p2]);
    int pt1 = x, pt2 = p[pt[x]];
    swap(pt[pt1], pt[pt2]);
    --lb[cnt[x] + 1]; ++cnt[x];
    for (auto node : ask[u]) {
        if (lb[node.l] + node.k - 1 > n) ans[node.id] = -1;
        else ans[node.id] = p[lb[node.l] + node.k - 1];
    }
    for (auto v : G[u]) {
        if (v == fa) continue;
        DFS(v, u);
    }
    --cnt[x]; ++lb[cnt[x] + 1];
    swap(pt[pt1], pt[pt2]);
    swap(p[p1], p[p2]);
}

int main() {
    Read(T);
    while (T--) {
        Read(n);Read(q);
        lb[0] = 1;
        for (int i = 1;i <= n;++i) { G[i].clear(); p[i] = pt[i] = i; lb[i] = n + 1; cnt[i] = 0; }
        for (int i = 1;i <= n;++i) { Read(a[i]); ask[i].clear(); }
        for (int i = 2;i <= n;++i) { int u; Read(u); G[u].push_back(i); G[i].push_back(u); }
        for (int i = 1;i <= q;++i) {
            int v, l, k;
            Read(v);Read(l);Read(k);
            ask[v].push_back((Node) { i, l, k });
        }
        DFS(1, 0);
        for (int i = 1;i <= q;++i)
            printf("%d ", ans[i]);
        printf("\n");
    }

    return 0;
}