[CF1794E] Labeling the Tree with Distances 题解

[CF1794E] Labeling the Tree with Distances 题解

题目描述

给你一个树，边权为 \(1\)。给定 \(n-1\) 个数，你需要将这些数分配到 \(n-1\) 个节点上。

一个点 \(x\) 是好的，当且仅当存在一种分配方案，所有被分配数的点到 \(x\) 的最短路径长度等于其被分配的数。

求所有好点。

思路

从特殊情况出发，如果有 \(n\) 个数的话怎么做。

可以发现，如果我们以 \(u\) 为根，所有点的深度的集合与这 \(n\) 个数的集合相等的话，那么 \(u\) 就是一个好点，如果一个一个点比较的话很耗时，可以使用哈希来描述一个集合，这是我的哈希函数：

\[h_u = c_0\times base^1+c_1\times base^2+\dots+c_{n-1}\times base^{n} \]

其中 \(c_i\) 表示深度为 \(i\) 的点的个数。

如果 \(h_u = tar\)，\(tar\) 是目标集合，那么 \(u\) 是好点。

这个时候如果我们每次求一遍 \(h_i\)，时间复杂度是 \(O(n^2)\) 的，但是这种与深度相关的 DP 很明显可以换根，所以时间复杂度降为 \(O(n)\)。

但是这只是一个特殊情况，如果少了一个数怎么办，其实也很好理解，如果 \(tar\) 中少了一个数，它的哈希值就会减少 \(base^k\)，我们只需要判断 \(h_u - tar\) 是否是一个 \(base\) 的次幂即可。

时间复杂度：\(O(n)\)。

单哈希会被卡，双哈希有时候也会寄，可以用三哈希。

#include <iostream>
#include <algorithm>
#include <cstring>
#include <queue>
#include <unordered_map>
#define int long long
using namespace std;

const int N = 2e5 + 10;

int n, a[N], p[N], f[N], dep[N], ans[N], base, mod, f2[N], tar, b[N];
unordered_map<int, bool> h;
vector<int> g[N];
void dfs(int u, int fa) {
    f[u] = base;
    for(auto v : g[u]) {
        if(v == fa) continue;
        dfs(v, u);
        f[u] = (f[u] + f[v] * base % mod) % mod;
    }
}
void dfs2(int u, int fa) {
    if(h.count((f2[u] % mod - tar + mod) % mod)) ans[u] ++;
    for(auto v : g[u]) {
        if(v == fa) continue;
        f2[v] = f[v] + base * ((f2[u] - base * f[v] % mod + mod) % mod) % mod;
        f2[v] %= mod;
        dfs2(v, u);
    }
}

void work() {
    tar = 0, p[0] = 1;
    h.clear();
    h[1] = 1;
    for(int i = 1; i <= n; i ++) p[i] = p[i - 1] * base % mod, h[p[i]] = 1;
    for(int i = 0; i < n; i ++) tar = (tar + b[i] * p[i + 1] % mod) % mod;
    dfs(1, 0), f2[1] = f[1], dfs2(1, 0);
}

signed main() {
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> n;
    for(int i = 1; i < n; i ++) cin >> a[i], b[a[i]] ++;
    for(int i = 1, a, b; i < n; i ++) {
        cin >> a >> b;
        g[a].push_back(b), g[b].push_back(a);
    }
    base = 19260817, mod = 998244353, work();
    base = 19491001, mod = 1011451423, work();
    base = 19421221, mod = 1e9 + 7, work();
    int cnt = 0;
    for(int i = 1; i <= n; i ++) if(ans[i] == 3) cnt ++;
    cout << cnt << '\n';
    for(int i = 1; i <= n; i ++) if(ans[i] == 3) cout << i << ' ';
    return 0;
}