성대나라에는 각 도시별로 가뭄을 대비하기 위한 물탱크가 하나씩 존재한다. 이 물탱크들은 모두 연결되어있으며, 루트(성대나라의 수도)가 있는 트리의 형태를 가진다. 지금 성대나라는 물탱크의 물을 사용하여 가뭄을 버텨냈으나, 그 영향으로 모든 물탱크에 물이 비어버리고 말았다.
성대나라의 물관리 시스템은 다소 특수해서, 물은 항상 다음과 같은 방식으로 채워진다:
A번 도시에 물을 채우기로 했다면, 수도에서부터 A번 도시까지 잇는 직선 경로에 수도부터 차례대로 1L, 2L, ⋯이 채워져서 A번 도시에는 (수도부터 A번 도시까지의 도시 수) L 만큼 추가된다.
예를 들어, 아래 그림과 같이 물탱크가 연결되어 있을 때, "4번 도시에 물을 채운다"라고 하면, 1번 도시에 1L, 4번 도시에 2L의 물이 추가된다. 만약 "5번 도시에 물을 채운다"라고 하면 1번 도시에 1L, 2번 도시에 2L, 5번 도시에 3L의 물이 추가된다.
성대나라의 물탱크 관리 담당인 균관이는 어느 도시에 몇 리터의 물이 저장되어있는지 자신이 궁금해질 때마다 알기를 원한다. 균관이를 도와주는 프로그램을 만들어보자.
첫째 줄에 성대나라의 도시의 수 N (1 ≤ N ≤ 200,000)과 수도의 번호 C (1 ≤ C ≤ N)가 공백으로 구분되어 주어진다.
둘째 줄부터 N-1개의 줄에 연결되어있는 두 도시의 번호 쌍 x, y가 공백으로 구분되어 주어진다(1 ≤ x, y ≤ N, x ≠ y). 물탱크의 연결 형태는 트리 구조임이 보장된다. N+1번째 줄에 질의의 수 Q(1 ≤ Q ≤ 200,000)가 주어진다. N+2번째 줄부터 Q개의 줄에 질의가 들어온다. 질의는 다음과 같이 두 종류 중 하나로 주어진다:
1 A : A도시에 물을 채운다. 2 A : 현재 A도시에 얼마만큼의 물이 채워져 있는지 출력하라.
위 문제는 오일러 경로 테크닉으로 세그먼트 트리의 갱신이 구간으로 이루어지는 lazy 형태로 바꿀 수 있었다. 하지만 본 문제에서는 현재 노드와 그 조상 노드들의 갱신을 요구한다. 골치 아파지는 문제이다. lazy propagation은 그 자손 노드로의 하향식 전파를 전제로 하므로 본 문제를 풀기에는 적절치 않다. 또한 상향식으로 넘버링을 진행할 경우 트리의 특성상 여러 노드가 한 조상을 공유할 수 있으므로, 적절한 해법이 되지 못한다.
여기서 다시 관점을 바꾸어 보자. 현재 i번째 트리 노드가 세그먼트 트리로 치환되었을때의 인덱스가 n[i], 그리고 그 트리 노드의 최대 자손 인덱스를 n[j]라고 하자. 그러면 i번째 노드가 포함하는 자손 노드들의 집합은 [n[i], n[j]]에 속한다.
우리는 이 i번째 트리 노드를 자손으로 하는 모든 조상 노드들이 n[i]를 포함하는 범위를 지니고 있음을 알 수 있다. 이를테면 루트 노드의 세그먼트 트리에서의 범위는 모든 범위를 포괄할 것이다. 또한 i번째 트리 노드들을 제외한 자손 노드들의 범위는 [n[i], n[j]]노드 범위 내에 포함되어 있으며, n[i]를 포함하지 못한다. 즉 이 둘로부터 유추할 수 있는 사실은, n[i]값이 본인 및 조상 노드들에만 영향을 끼친다는 것이다.
즉 우리는 i번째 노드를 갱신 시 n[i]값만을 업데이트하면 된다. n[i]값을 갱신하면 나중에 상향 노드를 탐색할 때 자동으로 반영이 될 것이기 때문이다. 또한 검색할 땐 n[i]에 간접적으로 상향으로 영향을 끼칠 수 있는 모든 값, 즉 n[i]의 모든 자손 값을 참조해야 하므로 [n[i], n[j]]구간합을 반영하면 된다. 정리하자면
갱신 시 n[i]만
검색 시 sum([n[i], n[j]])
가 되겠다. 여기서 i번 도시에 채워지는 물은 항상 그 도시의 depth만큼 채워지므로, 앞서 세그먼트 트리로 구한 값에 depth값만큼을 곱해주면 쿼리에 대한 정답을 구할 수 있다.
풀이 코드
from collections import defaultdict
import sys
sys.setrecursionlimit(1000000)
input = sys.stdin.readline
N, C = map(int, input().split())
edge_dict = defaultdict(list)
visited = [False]*N
node_info = [[0]*3 for _ in range(N)]
idx = -1
for _ in range(N-1) :
x, y = map(int, input().split())
edge_dict[x-1].append(y-1)
edge_dict[y-1].append(x-1)
def dfs(node, depth) :
global idx
visited[node] = True
idx += 1
node_info[node][0] = idx
node_info[node][2] = depth
for child in edge_dict[node] :
if not visited[child] :
dfs(child, depth+1)
node_info[node][1] = idx
class SegTree() :
def __init__(self) :
self.tree = [0]*(4*N)
def update(self, target) :
target = node_info[target][0]
def _update(start, end, idx) :
if target < start or target > end :
return
if start == end :
self.tree[idx] += 1
return
mid = (start + end) // 2
_update(start, mid, idx*2)
_update(mid+1, end, idx*2+1)
self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
_update(0, N-1, 1)
def search(self, target) :
left, right, mul = node_info[target]
def _search(start, end, idx) :
if right < start or left > end :
return 0
if left <= start <= end <= right :
return self.tree[idx]
mid = (start + end) // 2
lval = _search(start, mid, idx*2)
rval = _search(mid+1, end, idx*2+1)
return lval + rval
print(_search(0, N-1, 1) * mul)
tree = SegTree()
dfs(C-1, 1)
for _ in range(int(input())) :
q, c = map(int, input().split())
if q == 1 :
tree.update(c-1)
else :
tree.search(c-1)
