2020CCPC(長春) - Strange Memory(樹上啟發式合并+位運算)

題目大意：給出一棵 n 個點組成的有根樹，一號節點是根節點，現在要求實作 n * n 的公式：

題目分析：樹上啟發式合并，需要修改部分內部實作，如果可以想到樹啟的話，那么應該往子樹上去靠攏，當每個點作為子樹的根時，其可以作為 lca 然后去統計子樹中可以匹配的 ( u , v ) 點對，這個題目因為 a[ i ] ！= 0，換句話說，點 u , v , lca( u , v ) 一定是互不相同的三個點，極大程度上簡化了題目（因為三個點的形式一定是一個分叉的形狀，不可能是鏈狀的），換句話說，當 lca 確定后，可以列舉每個點作為點 u，然后去統計“除了點 u 到點 lca 的這條路徑上的點之外，以 lca 為根節點的子樹中，所有 a[ v ] = a[ u ] ^ a[ lca ] 的點 v，然后對 v^u 加和”，但如果每個點都去列舉的話，那么時間復雜度將會是 n * n 級別的，又因為如果每個點都列舉一遍的話，換句話說對于一個有貢獻的 ( u , v ) 點對來說，點 u 和點 v 都會被列舉一遍從而形成了浪費

再考慮樹啟是關于輕重鏈剖分的，然后訪問輕鏈的次數是 O( logn ) 級別的，因此可以只列舉輕鏈上的點的貢獻，然后去尋找滿足條件的點與其匹配即可，這樣時間復雜度就下降到了 nlogn，具體實作就是，對于某一條輕鏈來說，先維護其貢獻，然后再將其下標更新到資料結構中維護

下面講該如何實作，首先上面簡化的模型已經足夠可以完成公式，但本題實際上是要求記錄下標的異或和，比賽時沒有多想，直接對每個權值維護了一個 set，用于記錄當前權值下有多少個下標，維護答案的時候暴力列舉即可，本以為時間復雜度是 nlog^2n 的，TLE 的原因是因為 set 的那一層 log，于是想到每次增加和洗掉都是連續的一段，將 set 換成了 vector 便輕松將時間復雜度控制到了 nlogn 級別，交上去也真的 AC 了，賽后和潘學長還有zx學長討論過后才發現，自己實際上寫了個假演算法，，因為如果暴力去維護下標的話，時間復雜度其實是 O( k * n * logn )，這里的 k 是，整棵樹中出現次數最多的數字的出現次數，如果整棵樹都是同一個數字的話，那么時間復雜度將退化成 O( n * n * logn ) 級別的，，然鵝感謝出題人，卑微菜雞在這里給你磕頭了，咚咚咚，多謝出題人沒有刻意去卡這種資料，導致在隨機資料下，常數 k 好像很小很小，甚至比 std 跑的還快，，

然后講一下正解吧，因為我們需要維護的是下標的異或和，而對于異或而言，非常重要的一個性質就是，拆位之后每一位都相互獨立，所以我們不妨對于下標的每一位都單獨跑一次樹啟，最后將答案加和即可，對于某一位來說，假設 u == 1，對于匹配到的 v 來說，只有 v == 0 的位置才具有貢獻，對于 u == 0 而言，同理只有 v == 1 才有貢獻，所以類比于上一段的思路，對于每個權值 val 來說，上一段的做法維護的是一個 vector，里面儲存著 a[ x ] == val 的 x，也就是下標，而本段的思路是，對于每個權值 val 來說，記錄一下所有下標，在二進制下第 i 位中共出現了多少個 0 和多少個 1，只是換了一下實作思路而已，這樣實作的時間復雜度是嚴格 O( nlogn * 20 ) 的，20 的意思是需要將下標拆成 20 位，因為 2^20 > 1e6 >= a[ i ]

代碼：

//#pragma GCC optimize(2)
//#pragma GCC optimize("Ofast","inline","-ffast-math")
//#pragma GCC target("avx,sse2,sse3,sse4,mmx")
#include<iostream>
#include<cstdio>
#include<string>
#include<ctime>
#include<cmath>
#include<cstring>
#include<algorithm>
#include<stack>
#include<climits>
#include<queue>
#include<map>
#include<set>
#include<sstream>
#include<cassert>
#include<bitset>
using namespace std;

typedef long long LL;

typedef unsigned long long ull;

const int inf=0x3f3f3f3f;

const int N=1e5+100;

vector<int>node[N],temp;

bool vis[N];

int deep[N],num[N],son[N],a[N];

int cnt[(1<<20)+100][25][2];//拆位

LL ans;

void update(int id,int val)
{
	for(int i=0;i<20;i++)
		cnt[a[id]][i][(id>>i)&1]+=val;
}

LL search(int id,int num)
{
	LL ans=0;
	for(int i=0;i<20;i++)
		ans+=cnt[num][i][!((id>>i)&1)]*(1<<i);
	return ans;
}
 
void dfs_son(int u,int fa,int dep)
{
	deep[u]=dep;
	son[u]=-1;
	num[u]=1;
	for(auto v:node[u])
	{
		if(v==fa)
			continue;
		dfs_son(v,u,dep+1);
		num[u]+=num[v];
		if(son[u]==-1||num[v]>num[son[u]])
			son[u]=v;
	}
}

void cal(int u,int fa,int lca)
{
	temp.push_back(u);
	ans+=search(u,a[u]^a[lca]);
	for(auto v:node[u])
	{
		if(v==fa||vis[v])
			continue;
		cal(v,u,lca);
	}
}

void del(int u,int fa)
{
	update(u,-1);
	for(auto v:node[u])
	{
		if(v==fa||vis[v])
			continue;
		del(v,u);
	}
}
 
void dfs(int u,int fa,int keep)
{
	for(auto v:node[u])
	{
		if(v==fa||v==son[u])
			continue;
		dfs(v,u,0);
	}
	if(son[u]!=-1)
	{
		dfs(son[u],u,1);
		vis[son[u]]=true;
	}
	update(u,1);
	for(auto v:node[u])
	{
		if(v==fa||vis[v])
			continue;
		cal(v,u,u);
		for(auto it:temp)
			update(it,1);
		temp.clear();
	}
	if(son[u]!=-1)
		vis[son[u]]=false;
	if(!keep)
		del(u,fa);
}

int main()
{
#ifndef ONLINE_JUDGE
//  freopen("data.in.txt","r",stdin);
//  freopen("data.out.txt","w",stdout);
#endif
//  ios::sync_with_stdio(false);
	int n;
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		scanf("%d",a+i);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		node[u].push_back(v);
		node[v].push_back(u);
	}
	dfs_son(1,-1,0);
	dfs(1,-1,1);
	printf("%lld\n",ans);













    return 0;
}