DFS序+BIT离线统计子树中小于该节点的数量+三元计数

#include <bits/stdc++.h>

using namespace std;
#define ll long long
#define mem(a, b) memset(a,b,sizeof(a))
#define INF 0x3f3f3f3f
#define p_queue priority_queue
#define DBG printf("this is a input\n")
#define fi first
#define se second
#define mk(a, b) make_pair(a,b)
#define pb(a) push_back(a)

int T;
int n ;
int p[5];
int head[200005], cnt;
struct e{
    int t, next;
}edge[400005];
vector<pair<int,int> > son[100005];
struct BIT
{
    int tree[200005] ; //开 1 倍空间
    void init()
    {
        mem(tree,0) ;
    }
    int lowbit(int k)
    {
        return k & -k;
    }
    void add(int x , int k)  // a[x] += k
    {
        while(x <= n)  //维护的是 [1 , n] 的序列
        {
            tree[x] += k ;
            x += lowbit(x) ;
        }
    }
    int sum(int x)  // sum[l , r] = sum(r) - sum(l - 1)
    {
        int ans = 0 ;
        while(x != 0)
        {
            ans += tree[x] ;
            x -= lowbit(x) ;
        }
        return ans ;
    }
    int query(int l , int r)
    {
        return sum(r) - sum(l - 1) ;
    }
} bit ;
//  子树中大于当前节点  子树中小于当前节点
int minn[200005];
int dfn[200005], now = 0, size[200005], pa[200005];
ll ans = 0;
void dfs2()
{
    for (int i = 1 ; i <= n ; i ++)
    {
        //1 left , 2 right
        ll xi = 0 , yi = 0 , xo = 0 , yo = 0;
        for (int j = head[i] ; j != - 1 ; j = edge[j].next)
        {
            int v = edge[j].t;
            if(v != pa[i])
            {
                int x = minn[v];
                if (p[2] == 1) ans += 1ll * yi * (size[v] - x);
                else if (p[2] == 2) ans += 1ll * xi * (size[v] - x) + 1ll * yi * x;
                else ans += 1ll * xi * x;
                xi += x;
                yi += size[v] - x;
            }
        }
        xo = i - 1 - xi;
        yo = n - i - yi;
        if(p[2] == 1) ans += 1ll * yi * yo;
        else if(p[2] == 2) ans += 1ll * xi * yo + 1ll * yi * xo;
        else ans += 1ll * xi * xo;
    }
}
void Init()
{
    ans = 0, now = 0;
    mem(dfn,0);
    mem(minn,0);
    mem(size,0);
    mem(pa,0);
    mem(head,-1);
    cnt = 0;
}
void add (int f, int t)
{
    edge[cnt].t = t;
    edge[cnt].next = head[f];
    head[f] = cnt ++;
}
void dfs1(int root , int fa)
{
    size[root] = 1;
    dfn[root] = ++ now;
    for (int i = head[root]; i != -1 ; i = edge[i].next)
    {
        int v = edge[i].t;
        if(v != fa)
        {
            pa[v] = root;
            dfs1(v, root);
            size[root] += size[v];
        }
    }
}
void Deal()
{
    for (int i = 1 ; i <= n ; i ++)
    {
        for (int j = head[i] ; j != -1 ; j = edge[j].next)
        {
            int v = edge[j].t;
            if(v != pa[i])
                minn[v] = bit.query(dfn[v], dfn[v] + size[v] - 1);
        }
        bit.add(dfn[i],1);
    }
    for (int i = 1 ; i <= n ; i ++)
        bit.tree[i] = 0 ;
}
int main(void)
{
    scanf("%d", &T);
    while (T --)
    {
        scanf("%d", &n);
        Init();
        for (int i = 1 ; i <= 3 ; i ++)
            scanf ("%d", &p[i]);
        for (int i = 1 ; i < n ; i ++)
        {
            int u , v ;
            scanf("%d %d",&u , &v);
            add (u, v);
            add (v, u);
        }
        dfs1(1,1);
        Deal();
        dfs2();
        printf("%lld\n",ans);
    }

}
/*
2
5
1 2 3
3 1
3 5
1 2
1 4
 */

题目链接：https://www.codechef.com/problems/CHGORAM