AtCoder ABC #369 : G - As far as possible

バグ取りに手間取った…。
https://atcoder.jp/contests/abc369/tasks/abc369_g

問題

N頂点の木を成す無向グラフが与えられる。
各辺には長さが設定されている。

整数Kが与えられるので、N頂点中K頂点を選ぶことを考える。
頂点1から、K個の頂点を通過して戻るような最短経路を考える。
その長さが最長となるようにK頂点を選びたい。
K=1～Nそれぞれにおいて、最短経路の最長値を答えよ。

解法

頂点1を根とするグラフで、各頂点において根から未選択の各頂点までの距離を考える。
K点選ぶ場合、その距離が最大となるものを選びたい。
なお、距離が最大のものを選んだあと、そこから根頂点までの辺の長さは、以後0になると考えればよい。

そこで、区間加算・区間最大値を求めるSegTreeを用いて、点をDFS順に並べておいて、根頂点から各点までの長さを管理しよう。
このSegTreeで区間最大値を求めることで、次に選ぶする点を求めることができる。
辺の長さが0になるのは、区間加算で再現できる。

int N;
vector<pair<int,int>> E[202020];
pair<int,int> P[202020];

int id,L[202020],R[202020],re[202020];
ll D[202020];

template<class V,int NV> class SegTree_3 {
public:
	vector<V> val;
	vector<pair<V,int>> ma;
	SegTree_3(){
		int i;
		val.resize(NV*2); ma.resize(NV*2);
		FOR(i,NV) ma[i+NV]={0,-i};
		for(i=NV-1;i>=1;i--) ma[i]=max(ma[2*i],ma[2*i+1]);
	};
	
	pair<V,int> getval(int x,int y,int l=0,int r=NV,int k=1) {
		if(r<=x || y<=l || y<=x) return {-1LL<<60,0};
		if(x<=l && r<=y) return ma[k];
		auto p=max(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));;
		p.first+=val[k];
		return p;
	}
	
	void update(int x,int y, V v,int l=0,int r=NV,int k=1) {
		if(l>=r||y<=x) return;
		if(x<=l && r<=y) {
			val[k]+=v;
			ma[k].first+=v;
		}
		else if(l < y && x < r) {
			update(x,y,v,l,(l+r)/2,k*2);
			update(x,y,v,(l+r)/2,r,k*2+1);
			ma[k]=max(ma[k*2],ma[k*2+1]);
			ma[k].first+=val[k];
		}
	}
};
SegTree_3<ll,1<<20> st;


void dfs(int cur,int pre) {
	L[cur]=id++;
	re[L[cur]]=cur;
	FORR2(e,c,E[cur]) if(e==pre) {
		P[cur]={e,c};
		D[cur]=D[e]+c;
	}
	FORR2(e,c,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-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,N) {
		st.update(L[i],L[i]+1,D[i]);
	}
	ll sum=0;
	FOR(i,N) {
		auto p=st.getval(0,N);
		sum+=p.first*2;
		int cur=re[-p.second];
		st.update(L[cur],L[cur]+1,-1);
		while(P[cur].second) {
			st.update(L[cur],R[cur],-P[cur].second);
			P[cur].second=0;
			cur=P[cur].first;
			
		}
		
		cout<<sum<<endl;
	}
	
	
}

まとめ

バグ見つけるのに20分近くかかってしまった…。