CF1957F2 Frequency Mismatch (Hard Version) 题解

题意：

给定一棵 \(n\) 个点的树，每个点 \(v\) 有颜色 \(a_v\)。

现在有 \(Q\) 次询问，每次给出两条路径 \(u_1\leadsto v_1\) 和 \(u_2\leadsto v_2\) 还有一个数 \(K\)，令 \(x_c\) 为颜色 \(c\) 在路径 \(u_1\leadsto v_1\) 上出现的次数，\(y_c\) 为颜色 \(c\) 在路径 \(u_2\leadsto v_2\) 上出现的次数，请列出 \(K\) 个颜色 \(c\) 满足 \(x_c\neq y_c\)，如果不满 \(K\) 个则全部列出。

数据范围：

\(1\le n\le 10^5,1\le a_i\le 10^5,1\le Q\le 10^5,1\le K\le 10\)。

经典的 sum hash ，不知道可以看这篇

考虑给每个点随机赋一个权值，然后在树上建立主席树，维护到根节点的链上所有权值的出现次数的哈希值。

根据 sum hash 的结论，哈希值相等等价于该区间中每个出现的数的出现次数都和另一个区间的相同

反过来，如果两个节点的哈希值不相同，那么我们要找的 \(c\) 就在这个区间里，直接暴力往下找即可。

这样每次询问我们直接在主席树上边走边查就好了，这一部分见代码中的 qry() 函数。

当然你其实也可以赋一个质数权值，只不过更容易被卡而已，因此我建议质数和随机数都写，稳妥。

在主席树上查询到一个节点的复杂度是 \(\mathcal{O}(\log n)\) ，查到 \(k\) 个答案以后就会退出

因此时间复杂度为 \(\mathcal{O}(qk\log n)\)

代码：

#include <bits/stdc++.h>
using namespace std;
// #define int long long
// #define INF 0x3f3f3f3f3f3f3f3f
typedef unsigned long long ull;
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define N ((int)(1e5 + 55))
#define ls(at) t[at].ls
#define rs(at) t[at].rs

const int m = 1e5 + 10;
ull pw1[N], pw2[N], pw3[N];
mt19937 rd(time(0)); vector<int> vec[N];
int n, K, tot, cnt, rt[N], lg[N], a[N], dep[N], ans[N], fa[18][N]; 
struct SegmentTree { int ls, rs; ull v1, v2, v3; } t[N * 32];
void push_up(int at)
{
    t[at].v1 = t[ls(at)].v1 + t[rs(at)].v1;
    t[at].v2 = t[ls(at)].v2 + t[rs(at)].v2;
    t[at].v3 = t[ls(at)].v3 + t[rs(at)].v3;
}
int update(int pre, int l, int r, int x)
{
    int at = ++tot; t[at] = t[pre];
    if(l == r) {
        t[at].v1 += pw1[l]; t[at].v2 += pw2[l]; t[at].v3 += pw3[l];
        return at;
    }
    int mid = (l + r) >> 1;
    if(x <= mid) ls(at) = update(ls(pre), l, mid, x);
    else rs(at) = update(rs(pre), mid + 1, r, x);
    push_up(at); return at;
}
void init()
{
    pw1[0] = pw2[0] = pw3[0] = 1; 
    rep(i, 2, n) lg[i] = lg[i >> 1] + 1;
    rep(i, 1, m) pw1[i] = pw1[i - 1] * 131;
    rep(i, 1, m) pw2[i] = pw2[i - 1] * 13131;
    rep(i, 1, m) pw3[i] = rd();
}
void dfs(int u, int f)
{
    fa[0][u] = f; dep[u] = dep[f] + 1;
    rt[u] = update(rt[f], 1, m, a[u]);
    rep(i, 1, lg[n]) fa[i][u] = fa[i - 1][fa[i - 1][u]];
    for(int v : vec[u]) if(v != f) dfs(v, u);
}
int lca(int x, int y)
{
    if(dep[x] < dep[y]) swap(x, y);
    for(int i = lg[n]; ~i; i--)
        if(dep[fa[i][x]] >= dep[y]) x = fa[i][x];
    if(x == y) return x;
    for(int i = lg[n]; ~i; i--)
        if(fa[i][x] != fa[i][y]) { x = fa[i][x], y = fa[i][y]; }
    return fa[0][x];
}
struct node
{
    int x, y, d, fd;
    node lc() const { return node{ls(x), ls(y), ls(d), ls(fd)}; }
    node rc() const { return node{rs(x), rs(y), rs(d), rs(fd)}; }
    ull v1() const { return t[x].v1 + t[y].v1 - t[d].v1 - t[fd].v1; }
    ull v2() const { return t[x].v2 + t[y].v2 - t[d].v2 - t[fd].v2; }
    ull v3() const { return t[x].v3 + t[y].v3 - t[d].v3 - t[fd].v3; }
    bool eq(const node &o) const { return o.v1() == v1() && o.v2() == v2() && o.v3() == v3(); }
};
void qry(const node A, const node B, int l, int r)
{
    if(cnt == K || A.eq(B)) return;
    if(l == r) { ans[++cnt] = l; return; }
    int mid = (l + r) >> 1;
    qry(A.lc(), B.lc(), l, mid); qry(A.rc(), B.rc(), mid + 1, r);
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n; init();
    for(int i = 1; i <= n; i++) cin >> a[i];
    for(int i = 1, u, v; i < n; i++)
    {
        cin >> u >> v;
        vec[u].push_back(v); vec[v].push_back(u);
    }
    dfs(1, 0); int qwq; cin >> qwq;
    for(int x1, y1, x2, y2; qwq--; )
    {
        cin >> x1 >> y1 >> x2 >> y2 >> K; cnt = 0;
        int d1 = lca(x1, y1), d2 = lca(x2, y2);
        const node A = {rt[x1], rt[y1], rt[d1], rt[fa[0][d1]]};
        const node B = {rt[x2], rt[y2], rt[d2], rt[fa[0][d2]]};
        qry(A, B, 1, m); cout << cnt << " \n"[!cnt];
        rep(i, 1, cnt) cout << ans[i] << " \n"[i == cnt];
    }
    return 0;
}