Hello 2024 : D. 01 Tree

年明けいきなり大当たりした回。
https://codeforces.com/contest/1919/problem/D

問題

N個の葉を持つ完全二分木を考える。
その際、各点の2つの子への辺は、片方が0、片方が1の重みをもつとする。

数列Aは、根からi番目の葉までのパスの重みの和がA[i]であることを意味する。
そのような数列を構築できる木は存在するか判定せよ。

解法

A[i]=0となる葉は高々1個しかありえない。
そうでない場合、A[i]に対し、A[j]=A[i]-1となるi未満の最大のjと、A[k]=A[i]-1となるiより大きな最小のkを考える。
もしA[j+1]...A[i]がいずれもA[i]以上であるか、A[i]....A[k-1]がいずれもA[i]以上であれば、条件を満たすグラフが構築できる。

int T,N;
int A[202020];

template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	static V const def=1<<20;
	V comp(V l,V r){ return min(l,r);};
	
	SegTree_1(){val=vector<V>(NV*2,def);};
	V getval(int x,int y,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) return def;
		if(x<=l && r<=y) return val[k];
		return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
	}
	void update(int entry, V v) {
		entry += NV;
		val[entry]=v;
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
};
SegTree_1<int,1<<20> st;
vector<int> C[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		vector<pair<int,int>> V;
		FOR(i,N) {
			cin>>A[i];
			st.update(i,A[i]);
			C[A[i]].push_back(i);
		}
		int ok=1;
		int n0=0;
		FOR(i,N) {
			if(A[i]==0) {
				n0++;
			}
			else {
				auto it=lower_bound(ALL(C[A[i]-1]),i);
				int del=0;
				if(it!=C[A[i]-1].end()&&st.getval(i,*it)>=A[i]) del=1;
				if(it!=C[A[i]-1].begin()&&st.getval(*prev(it)+1,i)>=A[i]) del=1;
				if(del==0) ok=0;
			}
		}
		
		if(n0!=1) ok=0;
		if(ok) {
			cout<<"YES"<<endl;
		}
		else {
			cout<<"NO"<<endl;
		}
		FOR(i,N) C[A[i]].clear();
	}
}

まとめ

これよく本番あっさりと解法にたどりついたな。