kmjp's blog

競技プログラミング参加記です

AtCoder ABC #337 (トヨタ自動車プログラミングコンテスト2024#1) : G - Tree Inversion

誤読を重ねてしまった。
https://atcoder.jp/contests/abc337/tasks/abc337_g

問題

1~Nのラベル付き頂点からなる、木を成す無向グラフが与えられる。
f(u)は、以下を満たす頂点対(v,w)の個数とする。

  • uとvのパス上にwがあり、かつw>v

各頂点vにおけるf(v)を求めよ。

解法

まず木の頂点をDFS訪問順に並べ、BITで以下を管理する。

  • 今頂点vを処理している場合、SubTree内にある、v未満のラベルを持つ頂点数
  • 各頂点wに対するf(w)の値

ラベル番号vの大きい順に処理する。
vのうち1つの辺の先にn個のvより小さなラベルの頂点があったとする。
その場合、その辺の先のSubTree以外の頂点wに対しては、f(w)がnだけ加算される。
頂点数の算出も、f(w)を範囲に対し加算することもBITで行える。

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<ll,20> num,sum;

int N;
const int ME=300001;
vector<int> E[ME];
int L[ME],R[ME],D[ME],P[ME],rev[ME],eid;
ll ret[202020];

void EulerTour(int cur,int pre=0,int d=0) {
	if(pre==-1) ZERO(L),ZERO(R),eid=0;
	rev[eid]=cur;
	P[cur]=pre;
	L[cur]=eid++;
	D[cur]=d;
	FORR(e,E[cur]) if(e!=pre) EulerTour(e,cur,d+1);
	R[cur]=eid;
}



void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,N) {
		num.add(i,1);
	}
	EulerTour(0,-1);
	
	for(i=N-1;i>=0;i--) {
		int tot=i;
		num.add(L[i],-1);
		FORR(e,E[i]) {
			if(L[e]>L[i]) {
				int tt=(num(R[e]-1)-num(L[e]-1));
				sum.add(0,tt);
				sum.add(L[e],-tt);
				sum.add(R[e],tt);
				sum.add(N,-tt);
			}
			else {
				int tt=tot-(num(R[i]-1)-num(L[i]-1));
				sum.add(L[i],tt);
				sum.add(R[i],-tt);
			}
		}
		
		
	}
	
	FOR(i,N) cout<<sum(L[i])<<" ";
	cout<<endl;
	
}

まとめ

誤読多いのもったいないな。