求条 E
  • 板块学术版
  • 楼主封禁用户
  • 当前回复1
  • 已保存回复1
  • 发布时间2024/2/24 08:18
  • 上次更新2024/2/24 11:24:34
查看原帖
求条 E
1073459
封禁用户楼主2024/2/24 08:18

淀粉质代码。

#include<bits/stdc++.h>
using namespace std;
#define il inline
#define int long long 
#define re register
#define inf 123456789
il int read()
{
    int x = 0, f = 1; char c = getchar();
    while(c < '0' || c > '9') { if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - 48, c = getchar();
    return x * f;
}
#define rep(i, s, t) for(re int i = s; i <= t; ++ i)
#define drep(i, s, t) for(re int i = t; i >= s; -- i)
#define Next(i, u) for(re int i = head[u]; i; i = e[i].next)
#define mem(k, p) memset(k, p, sizof(k))
#define maxn 200005
struct edge{int v, next;}e[maxn << 1];
int n, m, head[maxn], cnt, k, ans, Ans;
int c[maxn];
bool flag[maxn];
int dp[maxn], vis[maxn], sum, rt, siz[maxn], rev[maxn], tot;
void add(int u, int v)
{
    e[++ cnt] = (edge){v, head[u]}, head[u] = cnt;
    e[++ cnt] = (edge){u, head[v]}, head[v] = cnt;
}
void getrt(int u, int fr)
{
//	cout<<u<<endl;
    siz[u] = 1, dp[u] = 0;
    Next(i, u)
    {
        int v = e[i].v;
        if(v == fr || vis[v]) continue;
        getrt(v, u);
        siz[u] += siz[v];
		dp[u] = max(dp[u], siz[v]);
    }
    dp[u] = max(dp[u], sum - siz[u]);
    if(dp[u] < dp[rt]) rt = u;
}
void getdis(int u, int fr)
{
    rev[++ tot] = c[u];
    Next(i, u)
    {
        int v = e[i].v;
        if(v == fr || vis[v] || flag[c[v]]) continue;
        flag[c[v]]=1;
        getdis(v, u);
        flag[c[v]]=0;
    }
}
int doit(int u)
{
    tot = 0;getdis(u, 0);
    sort(rev + 1, rev + tot + 1);
    ans = 0, cnt = 0;
    for (int i = 1; i <= tot; ++i) {
    	if (rev[i] == c[u]) ++ans;
		else {
			ans += cnt;
			if (rev[i] != rev[i + 1]) cnt = 0;
			else ++cnt;
		} 
		flag[rev[i]]=0;
	}
    return ans>0?ans-1:0;
}
int doit2(int u)
{
    tot = 0;getdis(u, 0);
    sort(rev + 1, rev + tot + 1);
    ans = 0, cnt = 0;
    for (int i = 1; i <= tot; ++i) {
    	if (rev[i] == c[u]) continue;
		ans += cnt;
		if (rev[i] != rev[i + 1]) cnt = 0;
		else ++cnt;
		flag[rev[i]]=0;
	}
    return ans;
}
void solve(int u)
{
//	cout<<u<<endl;
    vis[u] = 1;
	Ans += doit(u);
    Next(i, u)
    {
        int v = e[i].v;
        if(vis[v]) continue;
        Ans-=doit2(v); 
        sum = siz[v], dp[0] = n, rt = 0;
        getrt(v, u), solve(rt);
    }
}
signed main()
{
	int T=read();
	while(T--){
		Ans=0;
		cnt=0;
	    n = read();
	    rep(i, 1, n) vis[i]=head[i]=dp[i]=flag[i]=0,siz[i]=1,c[i] = read();
	    rep(i, 1, n - 1){int u = read(), v = read(); add(u, v);}
	    rt=0;
//	    cerr<<"42343223\n";
	    dp[0] = sum = n, getrt(1, 0);
		solve(rt);
	    printf("%lld\n", Ans);
	    rep(i, 1, n) vis[i]=head[i]=dp[i]=flag[i]=0,siz[i]=1;
	    cnt=Ans=0;
	}
    return 0;
}
2024/2/24 08:18
加载中...