Codeforces #695 Div2 E. Distinctive Roots in a Tree

なんか本番やたら苦戦してるな。
https://codeforces.com/contest/1467/problem/E

問題

木を成す無向グラフが与えられる。
各点には、整数値が設定されている。

以下の条件を満たす頂点vは何通りか。

• vを根とする根付き木を考える。根から各葉までのパス上にある頂点で、同じ値が2回以上登場することはない。

解法

はじめにDFS順で頂点番号を付け替えておく。
以後、同じ値を持つ頂点集合ごとに処理する。
各頂点にある辺の先に、同じ値を持つ頂点がある場合、その頂点の先に根があると、パス上に同じ値が2つあるので不可となる。
不可の区間を、BITを使って高速に足し合わせていこう。

int N;
int A[202020];
vector<int> E[202020];
map<int,set<int>> M;
int id,L[202020],R[202020],re[202020];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,20> bt;

void dfs(int cur,int pre) {
	re[id]=cur;
	L[cur]=id++;
	M[A[cur]].insert(L[cur]);
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur);
	R[cur]=id;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i];
	}
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,0);
	
	
	FORR(m,M) {
		set<int> S=m.second;
		if(S.size()==1) continue;
		FORR(s,S) {
			int cur=re[s];
			bt.add(L[cur],1);
			bt.add(L[cur]+1,-1);
			
			if(*S.begin()<s || R[cur]<=*S.rbegin()) {
				bt.add(L[cur],1);
				bt.add(R[cur],-1);
			}
			FORR(e,E[cur]) if(L[e]>L[cur]) {
				auto it=S.lower_bound(L[e]);
				if(it!=S.end() && *it<R[e]) {
					bt.add(0,1);
					bt.add(L[e],-1);
					bt.add(R[e],1);
				}
			}
		}
	}
	int ret=N;
	FOR(i,N) if(bt(i)>0) ret--;
	cout<<ret<<endl;
	
	
}

まとめ

本番ちょっと詰めが甘くてACしきれなかった。