Project Presentation Gym - 102191I （LCA + dfs序）

题目

https://cn.vjudge.net/problem/Gym-102191I

题意

给你n个点的树，每个点有权值，问所有的权值的本身及祖先有多少点

思路

所有相同点深度之和相加就是答案，但是会有点重复计算，要求LCA，但两两求LCA就会超时，只需要按DFS序排序，和前一个求LCA就可以了

代码

#include <bits/stdc++.h>

using namespace std;
struct node
{
    int x,nx;
}edge[1000006];
int n,m;
vector<int> qw[1000006];
int head[1000006];
int fa[1000006][25],deep[1000006];
int dfsx[1000006];
int p;
bool cmp(node a,node b)
{
    if(a.nx == b.nx) return a.x <b.x;
    else return a.nx < b.nx;
}
void add(int u,int v)
{
    edge[p].x = v;
    edge[p].nx = head[u];
    head[u] = p++;
}
void dfs(int u,int f)
{
    dfsx[u] = ++p;
    fa[u][0] = f;
    deep[u] = deep[f]+1;
    for(int i = head[u];i != -1;i = edge[i].nx)
    {
        dfs(edge[i].x,u);
    }
}
void init()
{
    for(int i = 1;i < 25;i++)
        for(int j = 1;j <= n;j++)
            fa[j][i] = fa[fa[j][i-1]][i-1];
}
int LCA(int x,int y)
{
    int deepx = deep[x],deepy = deep[y];
    if(deepx < deepy)
    {
        swap(x,y);
        swap(deepx,deepy);
    }
    for(int i = 24;i >= 0;i--)
    {
        if(deep[fa[x][i]] >= deepy)
            x = fa[x][i];
    }
    if(x == y) return x;
    for(int i = 24;i >= 0;i--)
    {
        if(fa[x][i] != fa[y][i])
            x = fa[x][i],y = fa[y][i];
    }
    return fa[x][0];
}
int main()
{
    scanf("%d%d",&n,&m);
    for(int i = 1;i <= n;i++)
    {
        int x;
        scanf("%d",&x);
        qw[x].push_back(i);
    }
    int rt;
    memset(head,-1,sizeof(head));
    memset(fa,0,sizeof(fa));
    p = 0;
    for(int i = 1;i <= n;i++)
    {
        int x;
        scanf("%d",&x);
        if(x == 0)
        {
            rt = i;
            continue;
        }
        add(x,i);
    }
    p = 0;
    dfs(rt,0);
    init();
    for(int i = 1;i <= m;i++)
    {
        int ans = 0;
        int len = qw[i].size();
        vector<node> q;
        for(int  j = 0;j < len;j++)
        {
            node a;
            a.x = qw[i][j];
            a.nx = dfsx[qw[i][j]];
            q.push_back(a);
        }
        sort(q.begin(),q.end(),cmp);
        if(len!=0)
            ans += deep[q[0].x];
        for(int j = 1;j < len;j++)
        {
            int u = q[j-1].x,v = q[j].x,lca = LCA(u,v);
            ans += deep[v]-deep[lca];
        }
        printf("%d ",ans);
    }
    return 0;
}