[Easy] CF1923E Count Paths

虚树 + 简单换根 dp。

简要题意

给定一个 $n$ 个点的树，点 $i$ 有点权 $c_i$，计算满足下列所有条件的路径 $(u,v)$ 数量：

• $u\neq v$。
• $c_u=c_v$。
• $(u,v)$ 路径上除了端点的其他点（可能不存在）的点权均与 $a_u,a_v$ 不同。

$T$ 组数据。$1\leq \sum n\leq 2\times 10^5,1\leq c_i\leq n$。

思路

首先考虑如果如果 $c_i\leq 10$ 怎么做，那么不妨枚举每一个点权 $x$，计算路径端点点权为 $x$，路径其余点点权不为 $x$ 的数量，则为了方便讨论，可以令 $a_i=[c_i=x]$，这样点权只有 $0/1$ 了。

这个形式就可以换根 dp。设 $f(i)$ 表示以 $i$ 的子树中的一个点权为 $1$ 的点为端点，$i$ 为路径上的一点（不为另一个端点）的方案数。

则有：

$$f(u)=\begin{cases} 1& a_u=1\\ \sum_{v\in\mathrm{son}(u)} f(v)& a_u=0 \end{cases}$$

接着就是换根，假如当前根为 $u$，若 $a_u=1$，则 $u$ 可以为一个端点，此时答案贡献是 $\sum_{v\in\mathrm{son}(u)} f(v)$。换根到 $v$ 的时候只需要令 $f(u)\gets f(u)-f(v)+f(\mathrm{fa}(u))$（若 $a_u=0$）表示去除 $v$ 的贡献，添加父亲的贡献。

注意这样对于路径 $(u,v),(v,u)$ 各会计算一次，所以答案要减半。

回到本题，由于颜色很多，不可能每次对全树跑一遍换根 dp，所以可以改为对每个颜色的点建虚树，这些点 $a_i=1$，则其余的辅助点 $a_i=0$，然后正常跑换根 dp 即可，由于对 $k$ 个点建虚树的点数仍然是 $O(k)$，所以复杂度是正确的。

单组数据时间复杂度 $O(n\log n)$。

代码

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

const int N = 2e5 + 5;
int n, a[N], dfn[N], top[N], dep[N], siz[N], father[N], son[N];
vector<int> b[N];
vector<int> g[N], t[N];
using i64 = long long;

void dfs1(int u, int fa){
    dfn[u] = ++dfn[0], dep[u] = dep[fa] + 1, father[u] = fa, siz[u] = 1;
    for(int v : g[u]){
        if(v == fa) continue;
        dfs1(v, u), siz[u] += siz[v];
        if(siz[v] > siz[son[u]]) son[u] = v;
    }
}

void dfs2(int u, int fa){
    if(son[u]) top[son[u]] = top[u], dfs2(son[u], u);
    for(int v : g[u]){
        if(v == fa || v == son[u]) continue;
        top[v] = v, dfs2(v,u);
    }
}

int lca(int x, int y){
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        x = father[top[x]];
    }
    return dep[x] < dep[y] ? x : y;
}

i64 f[N];
bool is[N];

void dp1(int u){
    for(int v : t[u]) dp1(v);
    if(is[u]) f[u] = 1;
    else{
        for(int v : t[u]) f[u] += f[v];
    }
}

i64 ans;

void dp2(int u, int fa){
    if(is[u]){
        for(int v : t[u]) ans += f[v];
        ans += f[fa];
    }
    i64 tmp = f[u];
    for(int v : t[u]){
        if(!is[u]) f[u] = tmp - f[v] + f[fa];
        dp2(v, u);
        f[u] = tmp;
    }
}

void solve(){
    cin >> n;
    for(int i=1;i<=n;i++) cin >> a[i], b[a[i]].push_back(i);
    for(int i=1,u,v;i<n;i++){
        cin >> u >> v;
        g[u].push_back(v), g[v].push_back(u);
    }
    top[1] = 1, dfs1(1, 0), dfs2(1, 0);
    for(int i=1;i<=n;i++){
        vector<int> vct; vct.swap(b[i]);
        if(vct.size() <= 1) continue;
        int siz = vct.size();
        for(int i : vct) is[i] = true;
        sort(vct.begin(), vct.end(), [](int x, int y){
            return dfn[x] < dfn[y];
        });
        for(int j=1;j<siz;j++) vct.push_back(lca(vct[j - 1], vct[j]));
        sort(vct.begin(), vct.end(), [](int x, int y){
            return dfn[x] < dfn[y];
        });
        vct.erase(unique(vct.begin(), vct.end()), vct.end());
        for(size_t j=1;j<vct.size();j++) t[lca(vct[j - 1], vct[j])].push_back(vct[j]);
        dp1(vct[0]), dp2(vct[0], 0);
        for(int j : vct) is[j] = false, f[j] = 0;
        for(size_t j=1;j<vct.size();j++){
            int x = lca(vct[j - 1], vct[j]);
            f[x] = 0, t[x].clear();
        }
    }
    cout << (ans >> 1) << '\n';
}

void clear(){
    for(int i=1;i<=n;i++) g[i].clear(), siz[i] = son[i] = 0;
    ans = 0;
}

signed main(){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    int t; cin >> t;
    while(t--) solve(), clear();
    return 0;
}

// Written by xiezheyuan