yukicoder : No.1769 Don't Stop the Game

なんとか解けてよかった。
https://yukicoder.me/problems/no/1769

問題

辺に重みが付いた木を成す無向グラフを考える。
頂点対(i,j)のうち、頂点iから頂点jに辺に沿って最短経路で移動する間、通った辺の重みのxorが常に1以上であるようなものは何通りか。

解法

X[v] := 0番の頂点からv番の頂点に移動する際の通った辺の重みのxor
とする。条件を満たす(i,j)は、iからjのパス上に、i意外にX[i]=X[v]となる頂点vが存在しないことである。

X[i]が一致するようなiごとに考える。
その際、頂点数の多い場合と少ない場合について分けて考えよう。

• X[i]が一致するiが多い場合
  • X[i]=xとする。元のグラフから、X[v]!=xであるような頂点同士のみを結ぶ辺のみを残した部分グラフを考え、それぞれ連結成分を求めておく。
  • iに対して、iに隣接する頂点cに対し、X[i]!=X[c]であるcの連結成分中の点がjとなりうる。
• X[i]が一致するiが少ない場合
  • X[i]=X[l]となるi以外の頂点lを考える。先にDFS順で頂点番号を振りなおしておくと、iから辺をたどっていく際lを通らないと到達できない頂点番号は、ある区間となる。
  • lごとにその区間の和集合を考えれば、それがiに対し条件を満たさないjとなる。

int N;
vector<pair<int,int>> E[202020];
int X[202020];
int L[202020],R[202020];
int id;
map<int,vector<int>> V;

int P[21][200005],D[200005];

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<202020> uf;

void dfs(int cur,int x) {
	X[cur]=x;
	V[x].push_back(cur);
	L[cur]=id++;
	FORR2(e,o,E[cur]) if(e!=P[0][cur]) D[e]=D[cur]+1, P[0][e]=cur, dfs(e,x^o);
	R[cur]=id;
}

int getpar(int cur,int up) {
	int i;
	FOR(i,20) if(up&(1<<i)) cur=P[i][cur];
	return cur;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y>>k;
		E[x-1].push_back({y-1,k});
		E[y-1].push_back({x-1,k});
	}
	dfs(0,0);
	FOR(i,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];

	ll ret=1LL*N*(N-1);
	
	FORR2(a,W,V) {
		if(W.size()<=500) {
			FORR(x,W) {
				vector<pair<int,int>> V;
				FORR(y,W) if(x!=y) {
					if(L[x]>L[y]&&L[x]<R[y]) {
						k=getpar(x,D[x]-D[y]-1);
						V.push_back({0,L[k]});
						V.push_back({R[k],N});
					}
					else {
						V.push_back({L[y],R[y]});
					}
				}
				sort(ALL(V));
				int r=0;
				FORR2(v1,v2,V) {
					if(v1<=r) {
						if(v2>r) {
							ret-=v2-r;
							r=v2;
						}
					}
					else {
						ret-=v2-v1;
						r=v2;
					}
				}
			}
		}
		else {
			uf.reinit();
			FOR(i,N) if(X[i]!=a) FORR2(e,c,E[i]) if(X[e]!=a) uf(i,e);
			FORR(x,W) {
				ret-=N-1;
				FORR2(e,c,E[x]) if(X[e]!=a) ret+=uf.count(e);
			}
		}
		
	}
	cout<<ret<<endl;
}

まとめ

X[i]が一致するものが多い場合と少ない場合で分けて考える、って方針は割とすぐ考えついたけど、それぞれ解法を定めるのに割と手間取った。