23836(어떤 우유의 배달목록 (Hard)) 풀이

이미지를 누르면 문제로 이동합니다

 

이 문제는 내용 자체는 간단하다.

트리가 처음에 주어지고, 쿼리로 2가지 입력 중 하나가 들어온다.

 

1번 쿼리는 트리의 정점  u에서 정점 v로 이동하는 경로에 u부터 0, 1, 2, 3 ... v에 k을 더하는 쿼리다.

2번 쿼리는 각 정점에 대해 얼마의 값이 더해졌는가, 그 합을 구하는 쿼리이다.

 

위 트리는 예제에서 제공하는 트리다.

예제에 나온대로 쿼리를 진행해보자.

 

일단, 맨처음 1 3 5 쿼리가 주어진다. 3에서부터 5까지 0부터 k까지를 더하는 것이다.

1번 쿼리를 진행하면 다음과 같이 될 것이다.

 

2번 쿼리는 1 4 5 이다. 이도 맞찬가지로 진행해주면,

다음과 같이 만들어진다는 것을 알 수 있다.

4에서부터 시작했기 때문에, 4에는 0, 3에는 1, 2에는 2, 5에는 3을 더해준 모습이다.

 

다음 쿼리는 1 1 2 이다.

이를 진행하면 다음과 같이 된다.

 

이 이후에는 2 2, 2 5 로 2번과 5번 정점의 점수를 묻는 쿼리다.

그러므로, 순서대로 4, 5 를 출력해주면 된다.

 

 

이제 문제를 풀어보자.

 

이 문제는 입력의 범위가 크기 때문에 단순한 나이브로 코드를 작성한다면 시간초과가 나올 것이다.

 

총 쿼리가 10만번이니까 시간복잡도는 많아봐야 QlogN, Q(logN)^2 정도가 될것이다.

 

그리고 1번 쿼리가 구간에 대해 업데이트를 진행하기 때문에 레이지 세그먼트 트리가 필요할 것으로 보인다.

 

즉 이 문제에선 트리에서의 임의의 구간, 레이지 세그먼트 트리를 다루는 알고리즘이 필요하다.

 

본론부터 말하자면, HLD( Heavy Light Decomposition ) 이 이를 수행할 수 있는 알고리즘이다.

 

알고리즘에 대해서는 깊게 설명하지 않겠다. 워낙 구현량도 많고, 어려운 알고리즘 이기 때문에 나보다 더 잘설명한 kks227 블로그나 justicehui 블로그를 참고하길 바란다.

 

HLD 의 특징에 대해 설명하자면, 이 문제에서 주어지는 거처럼 트리에서 특정 정점에서 정점 사이에 대한 업데이트를 연속된 구간들로 나타낼 수 있으며, 이를 logN, 각 과정당 세그먼트 트리를 logN 으로 (logN)^2 으로 우리가 원하는 1번 쿼리를 수행할 수 있다.

 

이제 (logN)^2 로 구간 업데이트를 할 수 있다고 하니, 우리는 어떻게 구간에 연속되는 값을 더할 것인지만 생각하면 된다.

 

만약 특정 구간 [s,e] 에 1부터 차례로 더해준다고 할때, 구간 [s,e] 의 합은 얼마일까?

 

이는 1부터 e-s+1 까지의 합을 구하면 된다. 이는 수열의 합으로 구할 수 있다.

 

 

S_n 이 등차수열의 합이고, n은 n번째까지를 나타낸다. a는 등차수열의 초항, d는 공차를 의미한다.

 

우리가 최종적으로 구해야되는 것은 합이다.

우리는 구간의 길이 n을 알고 있으니, 초항과 공차만 알면 된다.

공차는 1이 아니냐 라는 소리를 할 수 있지만, 구간이 여러번 업데이트 된다면 공차가 1이 아닐수도 있다.

 

예를 들어 [s,e] 구간에 1부터 시작하는 업데이트와 5부터 시작하는 업데이트가 있다고 치자.

그러면 구간 [s,e] 에는 1,2,3,4... 와 5,6,7,8... 이 더해져야할 것이다.

이 둘을 각각 a_n, 과 b_n 이라고 하고 a_n+b_n을 구해보자.

6,8,10,12,.... 이 나오는 것을 알 수 있다.

 

이를 보면 등차수열은 유지되지만, 초항은 a_1 + b_1 이 되었고, 공차는 d_a + d_b 가 되었음을 알 수 있다.

이를 접목 시켜, 구간이 업데이트 될때마다 업데이트 횟수와 초항값들의 합을 저장해준다면 우리는 구간의 합을 구할 수 있게 된다.

 

레이지 세그먼트 트리에서 제공하는 propagate 에서 자손에게 값을 넘기는 부분은 구간이 [s,mid], [mid+1,e] 인 부모 등차수열이라고 할 수 있으니 잘 처리해주면 된다.

 

세그는 됐으니, HLD 부분을 생각해보자.

HLD 는 u와 v가 있다면 밑에서 부터 올라가 u와 v가 같은 그룹에서 연속구간을 가지게 될때까지 올라가는 형태로 진행한다.

 

우리는 특정 구간에 대해 초항 값을 알아야할 필요가 있으니, HLD 가 진행할 때 uHead ~ u, vHead ~ v 구간을 deque 에 저장해준뒤, 마지막에 u에서 부터 차례로 진행하면서 업데이트를 진행해주면 된다.

 

void update(ll u,ll v){ 
	u = dfsn[u]; 
	v = dfsn[v];
	deque<pair<ll,ll>>dqu,dqv;
	if(cn[u] != cn[v]){ 
		while(1){
			ll uHead = cHead[cn[u]], vHead = cHead[cn[v]];

			if(depth[uHead] > depth[vHead]){

				dqu.push_back(make_pair(uHead,u));

				u = parent[uHead]; 
			}else{
				dqv.push_back(make_pair(vHead,v));

				v = parent[vHead];
			}
			
			if(cn[u] == cn[v])break;
		}
	}
	if(u<v){
		dqv.push_back({u,v});
	}else{
		dqu.push_back({v,u});
	}
	ll V=0;
	while(!dqu.empty()){
		pair<ll,ll> T = dqu.front();
		dqu.pop_front();
		ST.update_d(1,1,n,T.first,T.second,V+T.second-T.first);
		V+=T.second-T.first+1;
	}
	while(!dqv.empty()){
		pair<ll,ll> T = dqv.back();
		dqv.pop_back();
		ST.update_u(1,1,n,T.first,T.second,V);
		V+=T.second-T.first+1;
	}
}

 

이는 위와 같이 구현하면 된다.

 

이렇게 HLD 와 세그를 구성해주고, 입력에 맞게 코드를 작성해주면 된다.

 

 

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll MAX = 100010; // N 
struct nd{
	ll v,w,i;
};
ll n,u,v,w;
struct ud{
	ll u,d,ucnt,dcnt;
};
ll func(ll x){
	if(x<=0)return 0;
	return x*(x+1)/2;
}
class SegTree{
	public:
	ll tree[4*MAX];
	ud lazy[4*MAX];
	ll A[MAX];
	void construct(ll n,ll s,ll e){
		lazy[n]={0,0,0,0};
		if(s==e){
			tree[n]=0;
			return;
		}
		ll mid=(s+e)/2;
		construct(n*2,s,mid);
		construct(n*2+1,mid+1,e);
		tree[n]=tree[n*2]+tree[n*2+1];
	}
	void propagate(ll n,ll s,ll e){
		if(lazy[n].ucnt!=0){
			if(s!=e){
				ll mmid=(s+e)/2;
				lazy[n*2].u+=lazy[n].u;
				lazy[n*2].ucnt+=lazy[n].ucnt;
				lazy[n*2+1].ucnt+=lazy[n].ucnt;
				lazy[n*2+1].u+=lazy[n].u+(mmid-s+1)*lazy[n].ucnt;
			}
			tree[n]+=(e-s+1)*(lazy[n].u*2+(e-s)*lazy[n].ucnt)/2;
			lazy[n].u=0;
			lazy[n].ucnt=0;
		}
		if(lazy[n].dcnt!=0){
			if(s!=e){
				ll mmid=(s+e)/2;
				lazy[n*2].d+=lazy[n].d;
				lazy[n*2].dcnt+=lazy[n].dcnt;
				lazy[n*2+1].dcnt+=lazy[n].dcnt;
				lazy[n*2+1].d+=lazy[n].d-(mmid-s+1)*lazy[n].dcnt;
			}
			tree[n]+=(e-s+1)*(lazy[n].d*2-(e-s)*lazy[n].dcnt)/2;
			lazy[n].d=0;
			lazy[n].dcnt=0;
		}
	}
	void update_u(ll n,ll s,ll e,ll le,ll ri,ll v){
		propagate(n,s,e);
		if(ri<s||e<le)return;
		if(le<=s&&e<=ri){
			lazy[n].u+=v+(s-le);
			lazy[n].ucnt++;
			propagate(n,s,e); 
			return;
		}
		ll mid=(s+e)/2;
		update_u(n*2,s,mid,le,ri,v);
		update_u(n*2+1,mid+1,e,le,ri,v);
		tree[n]=tree[n*2]+tree[n*2+1];
	}
	void update_d(ll n,ll s,ll e,ll le,ll ri,ll v){
		propagate(n,s,e);
		if(ri<s||e<le)return;
		if(le<=s&&e<=ri){
			lazy[n].d+=v-(s-le);
			lazy[n].dcnt++;
			propagate(n,s,e); 
			//cout << n << ' ' <<s << ' ' << v << "\n";
			return;
		}
		ll mid=(s+e)/2;
		update_d(n*2,s,mid,le,ri,v);
		update_d(n*2+1,mid+1,e,le,ri,v);
		tree[n]=tree[n*2]+tree[n*2+1];
	}
	ll Ma(ll n,ll s,ll e,ll le,ll ri){
		propagate(n,s,e);
		if(ri<s||e<le)return 0;
		if(le<=s&&e<=ri)return tree[n];
		ll mid=(s+e)/2;
		return Ma(n*2,s,mid,le,ri)+Ma(n*2+1,mid+1,e,le,ri);
	}
};

ll N, C, dcnt; 
	
ll tSize[MAX], dfsn[MAX]; 
vector<ll> child[MAX]; 

ll parent[MAX], depth[MAX],cn[MAX]; 

ll eVertex[MAX];
ll cHead[MAX], cTail[MAX]; 

SegTree ST;

vector<nd> adj[MAX]; 

void dfs1(ll curr,ll prev){
	tSize[curr] = 1;
	for(int j=0;j<adj[curr].size();j++){
		ll nxt = adj[curr][j].v, d = adj[curr][j].w, en = adj[curr][j].i; 
		if(nxt!=prev){
			dfs1(nxt,curr);
			child[curr].push_back(nxt);
			tSize[curr] += tSize[nxt];
		}

	}
}
void dfs2(ll curr,ll prev, ll currDepth){

	dcnt++;
	ll u = dfsn[curr] = dcnt; 
	cn[u] = C;
	depth[u] = currDepth;

	if(cHead[C] < 0)cHead[C] = u; 
	cTail[C] = u; 

	if(child[curr].empty()){
		C++; 
		return;
	}

	ll chained = child[curr][0];
	for(int i=1;i<child[curr].size();i++){
		ll nxt = child[curr][i];
		if(tSize[chained] < tSize[nxt]) chained = nxt;
	}
	dfs2(chained,curr,currDepth+1);
	parent[dfsn[chained]] = u;

	for(int i=0;i<child[curr].size();i++){
		ll nxt = child[curr][i];
		if(nxt != chained)dfs2(nxt,curr,currDepth+1); 

		ll v = dfsn[nxt];
		parent[v] = u; 
	}


}

void initialize(ll n){
	N = n;
	dfs1(1,-1); 
	C = dcnt = 0;
	fill(cHead,cHead+MAX, -1);
	fill(cTail,cTail+MAX, -1);
	parent[1]=-1; 

	dfs2(1,-1,0);

	for(int i=1;i<=N;i++){
		ll u = dfsn[i];
		for(int j=0;j<adj[i].size();j++){
			ll nxt = adj[i][j].v, d = adj[i][j].w, en = adj[i][j].i;
			ll v = dfsn[nxt];

			if(depth[u]<depth[v]){ 
				eVertex[en] = v;
				ST.A[v] = d; 
			}

		}
	}

	ST.construct(1,1,N);


}
void update(ll u,ll v){ 
	u = dfsn[u]; 
	v = dfsn[v];
	deque<pair<ll,ll>>dqu,dqv;
	if(cn[u] != cn[v]){ 
		while(1){
			ll uHead = cHead[cn[u]], vHead = cHead[cn[v]];

			if(depth[uHead] > depth[vHead]){

				dqu.push_back(make_pair(uHead,u));

				u = parent[uHead]; 
			}else{
				dqv.push_back(make_pair(vHead,v));

				v = parent[vHead];
			}
			
			if(cn[u] == cn[v])break;
		}
	}
	if(u<v){
		dqv.push_back({u,v});
	}else{
		dqu.push_back({v,u});
	}
	ll V=0;
	while(!dqu.empty()){
		pair<ll,ll> T = dqu.front();
		dqu.pop_front();
		ST.update_d(1,1,n,T.first,T.second,V+T.second-T.first);
		V+=T.second-T.first+1;
	}
	while(!dqv.empty()){
		pair<ll,ll> T = dqv.back();
		dqv.pop_back();
		ST.update_u(1,1,n,T.first,T.second,V);
		V+=T.second-T.first+1;
	}
}

ll m,q,c;
ll A,x;
int main(){
	ios_base::sync_with_stdio(NULL);
	cin.tie(0);
	cout.tie(0);
	cin >> n;
	for(int i=1;i<n;i++){
		cin >> u >> v;
		adj[u].push_back({v,0,i});
		adj[v].push_back({u,0,i});
	}
	initialize(n);
	cin >> q;
	for(int i=0;i<q;i++){
		cin >> A;
		if(A==1){
			cin >> u >> v;
		}else{
			cin >> x;
		}
	}
}