【题解】洛谷P5588 小猪佩奇爬树

题意

传送门

题解

这里是一种时空复杂度均为$O(n\log n)$的暴力做法。感觉这道题顺着思路想下去还是挺简单的。

题目的要求实际上是对于同种颜色的点，都必须在同一条链上。

那么我们自然想到把每一种颜色单独处理，先找到每种颜色中深度最大的那个点，显然如果这种颜色能符合一条链的条件，那么这个最深的点一定是链的一端，我们设它为$x$。

接着我们考虑另一端的位置。这里出现了三种情况（如果不存在这种颜色那么直接为0，下面不再讨论）：

1. 只有一个点是这种颜色，即另一端就是$x$。
2. 另一端是$x$的祖先。
3. 另一端不是$x$的祖先。

其实普遍状况是第三种（需要优先考虑），前两种都是符合题目要求的条件的，可以直接计算（具体见代码，第一种我通过$pre$数组预处理了）。

如何判断是否存在第三种情况呢？其实很简单，我们在这种颜色中寻找不是$x$的祖先的点中深度最大的点即可。

如果我们发现存在第三种情况，那么我们就要尝试验证它是否符合题目要求的条件。具体而言，设另一端为$y$，我们就是想要知道x->y这条链上是否包含了所有颜色与它们相同的点，也就是要求出一条链上有多少点是这种颜色的。

那么就可以想到使用树上差分的思想。我们求出每个点到根节点每种颜色分别有几个点。然后把$x$和$y$的答案相加再减去$lca$的答案即可（$lca$处特判是否是当前颜色）。

然后问题转化为求一个点到根节点的路径上包含的每种颜色的点的数量。我们可以想到使用主席树来维护。每个节点在其父亲的基础上进行修改，这样就可以在$n\log n$复杂度内完成统计。

代码

#include <bits/stdc++.h>
#define MAX 2000005
#define MAXM 25000005
#define ll long long
#define mid ((l+r)>>1)
using namespace std;

template<typename T>
void read(T &n){
    
    
    n = 0;
    T f = 1;
    char c = getchar();
    while(!isdigit(c) && c != '-') c = getchar();
    if(c == '-') f = -1, c = getchar();
    while(isdigit(c)) n = n*10+c-'0', c = getchar();
    n *= f;
}

template<typename T>
void write(T n){
    
    
    if(n < 0) putchar('-'), n = -n;
    if(n > 9) write(n/10);
    putchar(n%10+'0');
}

int n, cnt, tot;
int head[MAX], vet[MAX*2], Next[MAX*2];
int col[MAX], sz[MAX], dep[MAX], f[MAX][21];
vector<int> v[MAX];

void add(int x, int y){
    
    
    cnt++;
    Next[cnt] = head[x];
    head[x] = cnt;
    vet[cnt] = y;
}

int s[MAXM], lc[MAXM], rc[MAXM], rt[MAX];
void build(int &p, int l, int r){
    
    
    if(!p) p = ++tot;
    if(l == r) return;
    build(lc[p], l, mid);
    build(rc[p], mid+1, r);
}

void update(int &p, int l, int r, int x, int last){
    
    
    p = ++tot;
    lc[p] = lc[last], rc[p] = rc[last], s[p] = s[last]+1;
    if(l == r) return;
    if(mid >= x) update(lc[p], l, mid, x, lc[last]);
    else update(rc[p], mid+1, r, x, rc[last]);
}

int query(int l, int r, int p, int x, int y, int z){
    
          //z:lca
    if(l == r) return s[x]+s[y]-2*s[z];
    if(mid >= p) return query(l, mid, p, lc[x], lc[y], lc[z]);
    else return query(mid+1, r, p, rc[x], rc[y], rc[z]);
}

void dfs(int x, int fa){
    
    
    dep[x] = dep[fa]+1, sz[x] = 1;
    f[x][0] = fa;
    for(int i = 1; i <= 20; i++) f[x][i] = f[f[x][i-1]][i-1];
    update(rt[x], 1, n, col[x], rt[fa]);
    for(int i = head[x]; i; i = Next[i]){
    
    
        int v = vet[i];
        if(v == fa) continue;
        dfs(v, x);
        sz[x] += sz[v];
    }
}

int Lca(int x, int y){
    
    
    if(dep[x] < dep[y]) swap(x, y);
    for(int i = 20; i >= 0; i--){
    
    
        if(dep[f[x][i]] >= dep[y]) x = f[x][i];
    }
    if(x == y) return x;
    for(int i = 20; i >= 0; i--){
    
    
        if(f[x][i] != f[y][i]) x = f[x][i], y = f[y][i];
    }
    return f[x][0];
}
int get(int x, int fa){
    
    
    for(int i = 20; i >= 0; i--){
    
    
        if(dep[f[x][i]] > dep[fa]) x = f[x][i];
    }
    return x;
}

bool cmp(int a, int b){
    
    
    return dep[a] > dep[b];
}

ll pre[MAX];

int main()
{
    
    
    cin >> n;
    for(int i = 1; i <= n; i++){
    
    
        read(col[i]);
        v[col[i]].push_back(i);
    }
    int x, y, lca = 0;
    for(int i = 1; i < n; i++){
    
    
        read(x), read(y);
        add(x, y), add(y, x);
    }
    build(rt[0], 1, n);
    dfs(1, 0);
    for(int i = 1; i <= n; i++){
    
    
        for(int j = head[i]; j; j = Next[j]){
    
    
            int v = vet[j];
            if(v == f[i][0]) continue;
            pre[i] += (ll)sz[v]*(n-sz[v]-1);
        }
        pre[i] += (ll)(n-sz[i])*(sz[i]-1);
        pre[i] /= 2;
        pre[i] += n-1;
    }
    for(int i = 1; i <= n; i++){
    
    
        if(v[i].empty()){
    
    
            write((ll)n*(n-1)/2), puts("");
            continue;
        }
        sort(v[i].begin(), v[i].end(), cmp);
        x = v[i][0];
        if(v[i].size() == 1){
    
    
            write(pre[x]), puts("");
            continue;
        }
        y = v[i][v[i].size()-1];
        for(int j = 1; j < v[i].size(); j++){
    
    
            lca = Lca(x, v[i][j]);
            if(lca != v[i][j]){
    
    
                y = v[i][j];
                break;
            }
        }
        if(y != lca){
    
    
            int t = query(1, n, i, rt[x], rt[y], rt[lca]);
            if(col[lca] == i) t++;
            if(t == v[i].size()){
    
    
                write((ll)sz[x]*sz[y]), puts("");
            }
            else puts("0");
        }
        else{
    
    
            write((ll)sz[x]*(n-sz[get(x, lca)])), puts("");
        }
    }

    return 0;
}