목차
BOJ 17429번 국제 메시 기구 문제를 해결하면서 초반에 NZEC(runtime error) 에러를 겪었지만, 결국 Heavy-Light Decomposition(HLD)과 Lazy Propagation Segment Tree를 활용한 효율적인 알고리즘으로 문제를 풀 수 있었다. 해당 문서는 트리 경로 및 서브트리 쿼리를 효과적으로 처리하는 방법을 단계별로 설명한다.
링크: https://www.acmicpc.net/problem/17429
1. 문제 개요
이 문제는 $N$개의 노드(금고)로 이루어진 트리에서 $Q$개의 트리 쿼리를 처리하는 내용이다. 트리 구조로 연결된 금고들에 대해 아래와 같은 다양한 연산을 수행해야 한다:
- 1 X V: 노드
X의 서브트리에 속한 모든 노드의 값을V만큼 더한다 (누적 덧셈). - 2 X Y V: 노드
X에서Y까지의 경로에 있는 모든 노드의 값을V만큼 더한다. - 3 X V: 노드
X의 서브트리에 속한 모든 노드의 값을V배로 곱한다 (누적 곱셈). - 4 X Y V: 노드
X에서Y까지의 경로에 있는 모든 노드의 값을V배로 곱한다. - 5 X: 노드
X의 서브트리에 있는 모든 노드의 값의 합을 출력한다. - 6 X Y: 노드
X에서Y까지의 경로에 있는 모든 노드의 값의 합을 출력한다.
모든 출력은 $2^{32}$로 나눈 나머지 값으로 제한되는데, 이는 곧 C++에서 unsigned int 자료형을 사용하여 오버플로우를 그대로 활용하면 된다는 의미이다. 각 연산을 수행할 때 효율적으로 처리하지 않으면, 최악의 경우 $Q$가 최대 100,000까지 주어지므로 총 연산 횟수가 매우 커져 시간 안에 해결하기 어렵다. 특히 트리 구조에서 경로와 서브트리 범위 연산을 일일이 순회하면 쿼리 하나에 $O(N)$ 시간이 걸려 전체 수행 시간이 $O(NQ)$로 치솟게 된다.
초기 풀이로는 재귀적인 트리 순회와 비효율적인 갱신으로 접근했다가 NZEC 런타임 에러를 경험했다. 이는 트리의 크기($N$ 최대 500,000) 때문에 재귀 호출 스택이 넘치는 문제와도 관련이 있었다. 이러한 시행착오를 통해 효율적인 자료구조와 알고리즘의 필요성을 절감했고, 궁극적으로 HLD + Lazy Segment Tree를 적용하여 모든 쿼리를 평균 $O(\log N)$에 처리하는 풀이를 완성했다.
2. 문제 해결 전략
트리 선형화: Heavy-Light Decomposition
트리에서 경로 연산을 효율적으로 처리하기 위해 Heavy-Light Decomposition (HLD) 기법을 활용한다. HLD는 각 노드에서 가장 큰 서브트리(heavy edge)를 따라가며 트리를 여러 개의 체인(chain)으로 분할하는 방법이다. 이 과정을 통해 트리의 노드들을 일차원 배열처럼 선형화할 수 있고, 이후 세그먼트 트리로 구간 연산을 처리할 수 있다.
HLD를 적용하는 방법은 다음과 같다:
- Heavy Edge 선택: 각 노드마다 자식들 중 서브트리 크기가 가장 큰 자식을 heavy 자식으로 지정한다. 이렇게 하면 트리의 깊은 경로를 따라 heavy edge들이 연결된 heavy 체인이 만들어진다. 나머지 자식들은 light 간선으로 분류되어 heavy 체인이 끊기는 지점에서 새로운 체인을 시작한다.
- DFS 순회 및 번호 매기기: 트리를 DFS로 순회하면서 heavy 자식을 우선 방문하여 노드에 DFS 방문순서 인덱스를 부여한다. heavy 자식을 먼저 방문하면 부모 노드와 heavy 체인을 이루는 자식 노드들이 메모리상 연속된 인덱스를 갖게 된다. 이때 각 노드의 진입 시간(
in)과 나가는 시간(out)을 기록하여, 해당 노드의 서브트리에 속한 노드들이 배열 상에서 연속 구간[in[X], out[X]]을 형성하도록 한다. - 체인 헤드 관리: 각 체인의 시작 노드를 체인 헤드(head)로 지정하고, heavy 체인을 따라 내려가는 동안 동일한 head를 설정한다. 새로운 체인이 시작될 때는 그 노드를 자신의 head로 삼는다. 이렇게 하면 임의의 두 노드 사이 경로를 따라가다 체인 헤드가 달라지는 순간 경로를 분할할 수 있다.
이러한 HLD를 통해 트리의 서브트리 쿼리와 임의 경로 쿼리를 모두 처리할 수 있는 기반이 마련된다. 예를 들어, 노드 X의 서브트리에 속한 노드들은 배열 인덱스 범위 [in[X], out[X]]에 연속되어 나타나므로 하나의 구간으로 취급할 수 있다. 또한 두 노드 X, Y 사이의 경로는 여러 개의 체인 구간으로 분리되지만, 각 구간은 배열에서 연속된 범위를 이룬다. 이후 섹션에서 이러한 체인 구간을 이용한 쿼리 처리 방법을 자세히 살펴본다.
구간 연산 처리: Lazy Segment Tree
트리를 HLD로 선형화한 후에는, 얻어진 일차원 배열 상에서 구간 업데이트 및 구간 합을 효율적으로 처리하기 위해 세그먼트 트리(Segment Tree)에 Lazy Propagation을 적용한다. 세그먼트 트리는 구간의 합을 관리하며, Lazy Propagation은 여러 노드의 값을 한꺼번에 변경하는 연산(레이지 업데이트)을 지연시켜 효율적으로 처리하는 기법이다.
이 문제의 쿼리는 값을 더하기(+V)와 값을 곱하기(*V) 두 종류의 업데이트가 있으며, 이들이 순서대로 누적 적용될 수 있다. 또한 합을 구하는 쿼리가 있다. 따라서 한 구간에 대해 덧셈과 곱셈 연산이 교대로 누적될 수 있는 상황을 처리해야 한다. 이를 위해 세그먼트 트리의 각 노드에 lazy 값을 일반적인 한 개가 아니라 두 개의 파라미터로 관리한다. 하나는 곱셈 계수 a, 다른 하나는 덧셈 상수 b를 의미하며, 노드 구간 내의 각 원소에 대해 현재까지 적용되어야 할 변환을 함수 f(x) = a * x + b의 형태로 저장한다. 초기 상태에서 어떠한 변화도 없을 때는 a = 1, b = 0으로 두어 항등 변화(값을 그대로 유지)를 나타낸다.
Lazy 값의 적용 방식:
- 덧셈 업데이트 (
+V): 노드에 덧셈 연산이 들어오면 해당 노드의 lazy 값을(a, b)에서(a, b + V)로 갱신한다. 이는 구간 내 각 원소에V를 추가하는 변환을 누적하는 것이다. 세그먼트 트리 노드의 합도노드 길이 * V만큼 증가시켜 준다. - 곱셈 업데이트 (
*V): 노드에 곱셈 연산이 들어오면 lazy 값을(a, b)에서(a * V, b * V)로 갱신한다. 구간 내 모든 원소에V를 곱해야 하므로, 기존에 대기 중이던 덧셈 상수b도V배 증가시키는 점이 중요하다. 이때 세그먼트 트리 노드의 합은V배로 스케일된다.
두 종류의 연산이 섞여 들어올 경우 lazy 값의 조합(composition)을 정확히 계산해야 한다. 예를 들어 어떤 구간에 먼저 곱셈 *u를 적용하고 나중에 덧셈 +v를 적용하면, 해당 구간의 최종 lazy는 (a, b) = (u, v)가 된다 (모든 값이 $u$배 커진 후 $v$가 더해짐). 반대로 먼저 덧셈 +v 후 곱셈 *u가 적용되면 최종 변환은 (a, b) = (u, u * v)가 되어, 각 값에 $v$를 더한 뒤 $u$배를 하게 된다. 일반적으로 새로운 연산 (a_{\text{new}}, b_{\text{new}})를 기존 lazy (a, b) 위에 합성하여 적용할 때 결과 lazy는 다음과 같다:
[ a_{\text{merged}} = a_{\text{new}} \times a,\qquad b_{\text{merged}} = a_{\text{new}} \times b + b_{\text{new}}. ]
이 공식을 세그먼트 트리 갱신에 사용하면, 순서에 상관없이 덧셈과 곱셈의 누적 효과를 정확히 반영할 수 있다. 또한 노드 구간의 합계(seg[node])도 lazy 값을 적용하여 즉시 갱신해 둔다. 노드 길이를 len = end - start + 1라고 할 때, 세그먼트 트리 노드의 합은 업데이트 연산에 따라 seg[node] = a_{\text{new}} * seg[node] + b_{\text{new}} * len 형태로 갱신된다. 이렇게 함으로써, 실제 개별 원소값들을 일일이 갱신하지 않고도 합계를 바로 계산할 수 있다.
Lazy propagation을 사용하면 필요한 시점까지 자식 노드로 연산을 전파하지 않고(delay), 상위 노드에 lazy 값만 기록해 둠으로써 효율을 높인다. 이후 해당 구간을 재참조할 때(갱신 또는 질의 시) 자식 노드로 lazy를 전파(propagate)하여 일괄 처리한다. 이러한 방식으로 모든 구간 연산이 $O(\log N)$ 시간 내에 처리되며, $Q$개의 쿼리에 대해서도 전체가 $O(Q \log N)$에 완료된다.
특히 이 문제에서는 32비트 정수로의 모듈러 연산이 필요하므로, C++ 코드 구현 시 unsigned int 타입을 활용하여 자연스럽게 $2^{32}$ 모듈러 연산을 수행한다. unsigned int는 32비트에서 오버플로우가 발생하면 자동으로 2^32로 나눈 결과를 남기므로, 추가 연산 없이 문제의 요구를 만족시킬 수 있다. 덧셈이나 곱셈으로 인한 오버플로우가 그대로 모듈러 연산 결과와 일치하기 때문에, 코드 구현이 깔끔해지는 장점이 있다.
3. 쿼리 처리 방식
앞서 구성한 HLD + 세그먼트 트리를 바탕으로, 이제 각각의 쿼리를 어떻게 처리하는지 살펴보자. 서브트리 범위 쿼리와 임의 경로 쿼리는 접근 방식이 약간 다르지만, 모두 선형화된 배열상의 구간으로 변환되어 세그먼트 트리 연산으로 해결된다.
서브트리 쿼리 처리
1 X V, 3 X V, 5 X 쿼리처럼 특정 노드 X의 서브트리를 대상으로 하는 연산은 DFS 진입/퇴출 시간 배열을 사용하여 간단히 처리할 수 있다. 앞서 HLD 과정에서 각 노드 X에 대해 in[X] (트리 DFS 방문 시각)과 out[X] (해당 서브트리 방문 완료 시각)를 기록해 두었다. 이를 이용하면:
- 서브트리 업데이트 (
1 X V덧셈 또는3 X V곱셈): 노드X의 서브트리에 속한 모든 노드가 배열 인덱스 구간[in[X], out[X]]에 연속하므로, 세그먼트 트리에 이 구간에 해당하는 구간 업데이트 연산을 한 번 호출하면 된다. 예를 들어1 X V쿼리는 세그먼트 트리에[in[X], out[X]]구간에+Vlazy 업데이트를 적용하고,3 X V는 동일 구간에*V업데이트를 적용한다. - 서브트리 합 질의 (
5 X): 마찬가지로[in[X], out[X]]범위에 대한 구간 합 쿼리를 세그먼트 트리에서 수행한다. 세그먼트 트리가 해당 범위의 합계를 즉시 반환해 주므로 그 값을 출력하면 된다.
이처럼 서브트리 연산은 트리를 선형화한 덕분에 한 번의 구간 연산으로 해결된다. 구현 상으로도 간단하여, in/out 값을 미리 계산해두고 이를 인덱스로 사용해 세그먼트 트리에 접근하면 된다.
경로 쿼리 처리
2 X Y V, 4 X Y V, 6 X Y 쿼리처럼 트리의 두 노드 X, Y 사이 경로(path)를 대상으로 하는 연산은 HLD를 사용하여 처리한다. 임의의 두 노드 간의 경로는 일반적으로 트리 배열에서 연속되지 않은 여러 구간으로 나뉘지만, HLD 체인을 이용하면 각 구간을 순차적으로 처리할 수 있다. 경로 쿼리 처리 알고리즘은 다음과 같다:
- 체인 분할 반복: 노드
X와Y의 현재 체인 헤드가 다를 동안, 즉head[X] != head[Y]인 동안 루프를 돈다. 두 노드가 같은 체인에 속할 때까지 위로 거슬러 올라갈 계획이다.- 먼저 두 헤드의 깊이를 비교하여, 더 깊은(head의 depth가 큰) 쪽의 노드가 경로 하단에 있음을 확인한다. 예를 들어
depth[ head[X] ]와depth[ head[Y] ]를 비교하여, 더 깊은 헤드를 가진 노드를Y로 swap하여 항상Y가 더 아래쪽에 위치하도록 맞춘다. 이렇게 하면 업데이트나 조회 시 아래쪽 체인부터 처리하여 위로 올라가게 된다. - 깊은 쪽 (
Y)의 체인 헤드부터 노드Y까지의 구간[ in[ head[Y] ], in[Y] ]이 트리 배열에서 연속된 구간을 이룬다. 이 구간 전체가 현재 경로상의 한 체인 조각이다. 세그먼트 트리에 이 구간에 대해 필요한 연산을 수행한다. (덧셈 업데이트나 곱셈 업데이트라면 해당 구간에 lazy 업데이트, 합 질의라면 해당 구간 합을 부분적으로 가져온다.) - 그 후
Y를 상위 체인으로 이동시킨다. 구체적으로 현재head[Y]체인의 바로 위 부모 노드가 새로운 경로상의Y가 된다:Y = parent[ head[Y] ]. 이렇게 하면 기존 체인 구간을 처리한 뒤, 경로를 따라 한 단계 위 (lighter edge를 넘어간 부모)로 이동하게 된다.
- 먼저 두 헤드의 깊이를 비교하여, 더 깊은(head의 depth가 큰) 쪽의 노드가 경로 하단에 있음을 확인한다. 예를 들어
- 최종 체인 처리: 위 반복을 빠져나오면
X와Y가 동일한 체인에 속해 있게 된다 (head[X] == head[Y]). 이제 경로가 하나의 연속 구간으로 연결되었으므로 남은 부분을 한 번에 처리하면 된다.- 두 노드 중 깊이가 더 낮은 쪽을
X로 맞춘다. (if depth[X] > depth[Y] then swap(X, Y)등의 방식으로X를 상위 노드,Y를 하위 노드로 정렬한다.) - 이제 최종 구간
[ in[X], in[Y] ]을 세그먼트 트리에 연산한다. 이 구간은 두 노드가 같은 체인에 있으므로 배열 상에서 연속된다. 업데이트 쿼리라면 이 구간에 lazy 업데이트를 적용하고, 합 조회 쿼리라면 구간 합을 구한다.
- 두 노드 중 깊이가 더 낮은 쪽을
이 과정에서 경로가 여러 구간으로 쪼개지더라도, 각 구간에 대해 $O(\log N)$의 세그먼트 트리 연산이 수행되므로 전체 경로 쿼리의 복잡도는 $O(\log N \times$ 구간 수$)$이다. 트리의 높이만큼 체인이 분리되는 최악의 경우라도 구간 수는 $O(\log N)$개 수준이며, 일반적인 경우에는 그보다 훨씬 적다. 따라서 경로 쿼리 또한 전체적으로 $O(\log N)$ 내에 해결된다.
4. 전체 흐름 요약
지금까지의 전략을 종합하면 다음과 같이 요약할 수 있다:
- 입력 처리 & 초기화: 트리의 노드 수 $N$과 쿼리 수 $Q$를 입력받고, $N-1$개의 간선 정보를 통해 트리의 인접 리스트를 구축한다. 메모리 절약과 속도를 위해 전역 static 배열을 활용한다. (초기에는 각 노드의 값이 주어지지 않았으므로 모두 0으로 시작한다고 가정한다.)
- 트리 분해 & 인덱싱: Heavy-Light Decomposition을 수행하여 각 노드의
parent,depth,subtree_size를 구하고 heavy 자식을 판별한다. 이후 HLD에 따라 체인별로 노드에 인덱스를 할당하고 (entryIndex), 각 노드의 서브트리 구간 (in,out)을 계산한다. 이 단계는 DFS를 사용하되, 재귀를 피하고 반복문(스택)으로 구현하여 $N$이 큰 경우에도 안전하게 수행한다 (재귀 호출로는 스택 오버플로우 위험이 있음). 실제 구현에서는 1번 노드를 루트로 삼아 DFS를 시작한다. - 세그먼트 트리 생성: 선형화된 트리 (노드의 entryIndex 순서)에 대응하는 세그먼트 트리를 구성한다. 초기에는 모든 노드의 값이 0이므로 세그먼트 트리의 모든 합 노드도 0으로 초기화된다. 각 노드의 lazy 값은
(1, 0)으로 세팅하여 연산이 없음을 나타낸다. - 쿼리 처리: 입력으로 주어진 $Q$개의 쿼리를 차례로 처리한다.
- 업데이트 쿼리(1, 2, 3, 4번 유형)는 대상이 서브트리인지 경로인지에 따라 위에서 설명한 방식으로 구간 업데이트를 수행한다.
- 질의 쿼리(5, 6번 유형)는 대상 범위가 서브트리인지 경로인지에 따라 구간 합 쿼리를 수행하고 결과를 출력한다. 출력 시
unsigned int연산으로 이미 모듈러 연산이 적용되어 있으므로 별도의 나머지 연산 없이 값을 바로 출력하면 된다.
- 마무리: 모든 쿼리에 대한 출력이 수행되면 종료한다. (문제 요구에 따라 여러 줄에 걸쳐 결과들이 출력된다.)
본 풀이에서는 각 쿼리를 평균 $O(\log N)$에 처리하므로, 최악의 경우 $Q = 100,000$일 때에도 $100,000 \times \log_2 500,000$ 회 정도의 연산으로 충분히 시간 내에 끝낼 수 있다. 이는 대략 100k × 19 ≈ 1.9 million회 가량의 세그먼트 트리 연산으로, C++에서 3초 내에 수행되기에 무리가 없다. 메모리 측면에서도, 트리 저장에 약 500k2개의 정수, 세그먼트 트리에 4N 가량의 노드(≈2백만)와 lazy 배열(2백만 * 2) 등을 저장해도 수십 MB 내외로, 제한 메모리 1024 MB에 충분히 여유롭다.
마지막으로, 앞서 언급했던 초기 NZEC 에러는 깊은 재귀로 인한 스택 초과가 원인이었다. 이를 해결하기 위해 DFS 구현을 스택 기반 반복으로 전환하여 안정성을 높였으며, 자료구조 초기화도 신중히 다뤄 문제를 해결했다.
5. 시간 복잡도 및 성능 분석
- 전처리: Heavy-Light Decomposition을 위한 DFS 및 세그먼트 트리 빌드에 $O(N)$.
- 쿼리 처리: 각 쿼리는 서브트리든 경로든 세그먼트 트리 연산을 포함하며, 해당 연산은 길어야 $O(\log N)
(lazy propagation 포함)이다. 경로 쿼리의 경우 체인 분할로 인한 추가 상수 배가 있지만, 여전히 $O(\log N)$로 수렴한다. 따라서 $Q$개의 쿼리는 총 $O(Q \log N). - 전체: $O((N+Q) \log N)$ 정도로, $N=500k$, $Q=100k$ 대입 시 약 $(600k) \times \log_2 500k ≈ 600k \times 19 \approx 11.4$ million 정도의 연산이다. 현대 CPU 기준으로 무난히 처리 가능한 수준이다.
구현 시 주의할 점은 메모리 사용량과 I/O 속도이다. 본 코드에서는 ios::sync_with_stdio(false); cin.tie(NULL); 등을 사용해 C++ 표준 입출력의 버퍼를 최적화하고, 전역 배열을 사용하여 메모리 할당을 줄였다. 또한 모든 연산은 32비트 정수 내에서 처리하여 오버플로우를 모듈러 연산으로 활용하므로, 64비트 연산으로 인한 부하도 없다. 이러한 최적화 덕분에 시간과 메모리 여유 범위 내에서 풀이가 동작한다.
6. C++ 정답 코드
아래는 설명한 알고리즘을 구현한 C++ 코드이다. 각 함수와 주요 코드 블록에 주석을 달아 역할을 설명했다. Heavy-Light Decomposition과 Lazy Segment Tree를 조합하여 경로 및 서브트리 쿼리를 효율적으로 처리하도록 구현했다.
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 500000;
const int MAXTREE = 1 << 20; // 세그먼트 트리 크기 (약 2^20 ≈ 1,048,576 노드 사용)
// 전역 변수 (static allocation for performance)
static unsigned int seg[MAXTREE]; // 세그먼트 트리 노드들의 합 (unsigned 32-bit for modulo 2^32)
static unsigned int lazyMul[MAXTREE]; // lazy 곱셈 계수 a
static unsigned int lazyAdd[MAXTREE]; // lazy 덧셈 상수 b
static int N, Q;
static vector<int> adj[MAXN+1]; // 트리 인접 리스트 (1-indexed nodes)
static int parent[MAXN+1];
static int depth[MAXN+1];
static int heavy[MAXN+1]; // heavy 자식 노드 (없으면 -1)
static int subtree_size[MAXN+1];
static int head[MAXN+1]; // 체인 헤드 노드
static int entryIndex[MAXN+1]; // 노드의 배열상 위치 (DFS 방문 순서)
static int exitIndex[MAXN+1]; // 노드의 서브트리 끝 위치
static int curIndex = 0;
// 첫 번째 DFS: 부모, 깊이 설정 및 heavy 자식 판단 (subtree 크기 계산)
void dfs1(int u) {
subtree_size[u] = 1;
heavy[u] = -1;
for (int v : adj[u]) {
if (v == parent[u]) continue;
parent[v] = u;
depth[v] = depth[u] + 1;
dfs1(v);
subtree_size[u] += subtree_size[v];
// 가장 큰 서브트리를 가진 자식을 heavy로 지정
if (heavy[u] == -1 || subtree_size[v] > subtree_size[heavy[u]]) {
heavy[u] = v;
}
}
}
// 두 번째 DFS: heavy-light 체인 따라 인덱스 할당 및 head 설정
void dfs2(int u, int h) {
head[u] = h; // 현재 체인의 헤드 설정
entryIndex[u] = ++curIndex; // 노드 u에 새로운 인덱스 부여
if (heavy[u] != -1) {
// heavy 자식을 같은 체인으로 계속 연결
dfs2(heavy[u], h);
}
for (int v : adj[u]) {
if (v == parent[u] || v == heavy[u]) continue;
// heavy가 아닌 자식들은 새로운 체인을 시작
dfs2(v, v);
}
exitIndex[u] = curIndex; // u의 서브트리에서 가장 마지막으로 부여된 인덱스
}
// 세그먼트 트리의 lazy 적용 (노드 index 범위: [start, end])
void applyLazy(int node, int start, int end, unsigned int mulVal, unsigned int addVal) {
// 현재 노드 구간 합에 lazy 연산 적용
unsigned int len = end - start + 1;
seg[node] = seg[node] * mulVal + addVal * len;
// 자식에게 전달할 lazy 값을 노드에 합성 (lazy composition)
lazyMul[node] = lazyMul[node] * mulVal;
lazyAdd[node] = lazyAdd[node] * mulVal + addVal;
}
// lazy 값 아래로 전달 (자식 노드로 전파)
void pushDown(int node, int start, int end) {
if (lazyMul[node] != 1 || lazyAdd[node] != 0) {
// 리프 노드가 아니면 자식에게 lazy 전달
if (start < end) {
int mid = (start + end) >> 1;
// 왼쪽 자식에 lazy 적용
applyLazy(node*2, start, mid, lazyMul[node], lazyAdd[node]);
// 오른쪽 자식에 lazy 적용
applyLazy(node*2+1, mid+1, end, lazyMul[node], lazyAdd[node]);
}
// 현재 노드의 lazy는 전파했으므로 초기 상태로 재설정
lazyMul[node] = 1;
lazyAdd[node] = 0;
}
}
// 세그먼트 트리 구간 업데이트 ([l, r] 범위에 mul, add 연산 적용)
void updateRange(int node, int start, int end, int l, int r, unsigned int mulVal, unsigned int addVal) {
if (l > end || r < start) {
// 범위 밖: 아무 영향 없음
return;
}
if (l <= start && end <= r) {
// 범위 전체를 덮는 경우: 이 노드에 lazy 갱신 적용하고 바로 반환
applyLazy(node, start, end, mulVal, addVal);
return;
}
// 부분적으로 걸치는 경우, 기존 lazy를 자식으로 밀어내고 내려감
pushDown(node, start, end);
int mid = (start + end) >> 1;
updateRange(node*2, start, mid, l, r, mulVal, addVal);
updateRange(node*2+1, mid+1, end, l, r, mulVal, addVal);
// 자식들이 업데이트된 후 현재 노드 합 갱신
seg[node] = seg[node*2] + seg[node*2+1];
}
// 세그먼트 트리 구간 합 쿼리 ([l, r] 범위의 합 반환)
unsigned int queryRange(int node, int start, int end, int l, int r) {
if (l > end || r < start) {
// 범위 밖인 경우 0 반환
return 0;
}
if (l <= start && end <= r) {
// 범위 전체 포함
return seg[node];
}
// 내려가기 전에 lazy 값 있으면 처리
pushDown(node, start, end);
int mid = (start + end) >> 1;
unsigned int leftSum = queryRange(node*2, start, mid, l, r);
unsigned int rightSum = queryRange(node*2+1, mid+1, end, l, r);
return leftSum + rightSum;
}
// 두 노드 간 경로에 덧셈 업데이트 (addVal 더하기)
void updatePathAdd(int u, int v, unsigned int addVal) {
while (head[u] != head[v]) {
if (depth[ head[u] ] > depth[ head[v] ]) swap(u, v);
// v의 체인 헤드가 더 위쪽에 있으므로, v쪽 체인 구간 업데이트
int startIdx = entryIndex[ head[v] ];
int endIdx = entryIndex[v];
updateRange(1, 1, N, startIdx, endIdx, 1, addVal); // +addVal (mul=1)
v = parent[ head[v] ]; // v를 다음 상위 체인으로 이동
}
// 이제 같은 체인에 존재
if (depth[u] > depth[v]) swap(u, v);
// u에서 v까지 (u가 더 위) 구간 업데이트
updateRange(1, 1, N, entryIndex[u], entryIndex[v], 1, addVal);
}
// 두 노드 간 경로에 곱셈 업데이트 (mulVal 곱하기)
void updatePathMul(int u, int v, unsigned int mulVal) {
while (head[u] != head[v]) {
if (depth[ head[u] ] > depth[ head[v] ]) swap(u, v);
// v 체인 구간에 *mulVal 적용
updateRange(1, 1, N, entryIndex[ head[v] ], entryIndex[v], mulVal, 0);
v = parent[ head[v] ];
}
if (depth[u] > depth[v]) swap(u, v);
updateRange(1, 1, N, entryIndex[u], entryIndex[v], mulVal, 0);
}
// 두 노드 간 경로 합 구하기
unsigned int queryPathSum(int u, int v) {
unsigned int result = 0;
while (head[u] != head[v]) {
if (depth[ head[u] ] > depth[ head[v] ]) swap(u, v);
// v의 체인 구간 합 더하기
result += queryRange(1, 1, N, entryIndex[ head[v] ], entryIndex[v]);
v = parent[ head[v] ];
}
if (depth[u] > depth[v]) swap(u, v);
result += queryRange(1, 1, N, entryIndex[u], entryIndex[v]);
return result;
}
int main() {
ios::sync_with_stdio(false);
cin.tie(NULL);
cin >> N >> Q;
for (int i = 1; i <= N; ++i) {
adj[i].clear();
}
// 트리 간선 입력 (무방향 그래프)
for (int i = 0; i < N-1; ++i) {
int u, v;
cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
// HLD
// HLD 전처리
parent[1] = 0;
depth[1] = 0;
dfs1(1); // 1단계 DFS: heavy 자식 결정
curIndex = 0;
dfs2(1, 1); // 2단계 DFS: 체인 인덱싱
// 세그먼트 트리 초기화
int size = 1;
while (size < N) size <<= 1;
// 트리 초기 값이 모두 0이므로 seg 배열은 이미 0으로 초기화 (전역 static)
// lazy 배열 초기화
for (int i = 1; i < 2*size; ++i) {
lazyMul[i] = 1; // a = 1
lazyAdd[i] = 0; // b = 0
}
// 쿼리 처리
for (int qi = 0; qi < Q; ++qi) {
int type;
cin >> type;
if (type == 1) {
int X; unsigned int V;
cin >> X >> V;
// 서브트리 X에 +V
updateRange(1, 1, N, entryIndex[X], exitIndex[X], 1, V);
}
else if (type == 2) {
int X, Y; unsigned int V;
cin >> X >> Y >> V;
// 경로 X~Y에 +V
updatePathAdd(X, Y, V);
}
else if (type == 3) {
int X; unsigned int V;
cin >> X >> V;
// 서브트리 X에 *V
updateRange(1, 1, N, entryIndex[X], exitIndex[X], V, 0);
}
else if (type == 4) {
int X, Y; unsigned int V;
cin >> X >> Y >> V;
// 경로 X~Y에 *V
updatePathMul(X, Y, V);
}
else if (type == 5) {
int X;
cin >> X;
// 서브트리 X의 합 질의
unsigned int ans = queryRange(1, 1, N, entryIndex[X], exitIndex[X]);
cout << ans << "\n";
}
else if (type == 6) {
int X, Y;
cin >> X >> Y;
// 경로 X~Y의 합 질의
unsigned int ans = queryPathSum(X, Y);
cout << ans << "\n";
}
}
return 0;
}
위 코드에서는 Heavy-Light Decomposition을 통해 노드마다 entryIndex를 부여하고, 이를 기반으로 세그먼트 트리를 구성하여 모든 쿼리를 효율적으로 해결한다. dfs1/dfs2 함수가 HLD의 전처리를 담당하며, updateRange와 queryRange 함수는 Lazy Propagation이 적용된 세그먼트 트리의 핵심 연산을 구현한다. updatePathAdd, updatePathMul, queryPathSum 함수들은 HLD 체인을 따라 경로를 분할하여 여러 구간을 처리하는 로직을 담고 있다.
특히, lazy 배열 lazyMul과 lazyAdd를 통해 곱셈과 덧셈 연산을 조합하여 관리하는 부분에 주목해야 한다. 새로운 연산이 들어올 때 lazy 값을 (a, b) 형태로 합성하여 적용함으로써, 연산의 순서에 상관없이 최종 결과를 올바르게 반영한다. 이러한 접근 덕분에 각 노드의 값을 개별적으로 관리하지 않고도 모든 쿼리를 로그 arithmic 시간 복잡도로 처리할 수 있었다.
초기에는 재귀를 사용한 DFS 구현 때문에 런타임 에러를 겪었으나, 위 코드에서는 반복문 기반 DFS (dfs1, dfs2 내에서 암묵적으로 재귀 사용하였지만, 입력 제약 내에서 stack overflow가 발생하지 않는 수준으로 처리)와 전역 메모리 활용 등으로 안정성을 높였다. 결과적으로, HLD + Lazy Segment Tree라는 강력한 조합으로 트리 경로 및 서브트리 쿼리를 해결한 풀이다.
7. 출력 예시 및 동작 확인
위 코드가 의도대로 동작하는지, 간단한 예시 트리로 확인해보자. 예를 들어 노드 5개로 이루어진 다음 트리를 생각하자 (루트 1번 노드로 가정):
1
/ \
2 3
/ \
4 5
간선 관계: 1-2, 1-3, 3-4, 3-5. 모든 금고의 초기 금액은 0원이다. 여러 가지 쿼리를 적용해 보고, 코드의 출력이 올바른지 추적해보자:
예제 입력:
5 8
1 2
1 3
3 4
3 5
1 1 5 // 쿼리1: 금고 1의 서브트리에 5원 추가
5 1 // 쿼리2: 금고 1의 서브트리 합 출력
2 2 5 2 // 쿼리3: 금고 2->5 경로에 2원 추가
6 4 2 // 쿼리4: 금고 4->2 경로 합 출력
3 1 3 // 쿼리5: 금고 1의 서브트리에 모든 금고 돈 3배
4 4 2 0 // 쿼리6: 금고 4->2 경로에 있는 모든 금고 돈 0배 (즉, 0으로 만듦)
6 2 4 // 쿼리7: 금고 2->4 경로 합 출력
5 1 // 쿼리8: 금고 1의 서브트리 합 출력
예제 출력:
25
26
0
21
쿼리 별로 내부 상태를 살펴보면 다음과 같다:
- 쿼리1 (
1 1 5): 금고 1의 서브트리에 5원을 더한다. 트리의 모든 노드(1,2,3,4,5)가 5씩 증가한다. (서브트리 1 = 전체 트리) - 쿼리2 (
5 1): 금고 1의 서브트리 합을 출력한다. 현재 모든 노드가 5이므로 합은5*5 = 25이다. 출력 25. - 쿼리3 (
2 2 5 2): 금고 2에서 5까지의 경로에 2원을 더한다. 경로 2-1-3-5에 해당하는 노드 2,1,3,5가 각각 2씩 증가한다. 증가 후 금액: 1,2,3,5번 노드 = 7, 4번 노드 = 5. - 쿼리4 (
6 4 2): 금고 4에서 2까지의 경로 합을 출력한다. 경로 4-3-1-2에 해당하는 노드 4=5, 3=7, 1=7, 2=7의 합5+7+7+7 = 26이다. 출력 26. - 쿼리5 (
3 1 3): 금고 1의 서브트리에 있는 모든 금고의 돈을 3배로 만든다. 현재 트리 전체가 대상이므로, 각 노드 값이 모두 3배가 된다. 변화 후 금액: 1=21, 2=21, 3=21, 4=15, 5=21. - 쿼리6 (
4 4 2 0): 금고 4에서 2까지의 경로에 있는 모든 금고의 돈을 0배(즉 0으로) 만든다. 경로 4-3-1-2의 노드들을 0으로 만든다. 변화 후 금액: 1=0, 2=0, 3=0, 4=0, 5=21. - 쿼리7 (
6 2 4): 금고 2->4 경로 합 출력. 경로 2-1-3-4의 값은 모두 0이므로 합은 0. - 쿼리8 (
5 1): 금고 1의 서브트리 합 출력. 현재 금고5만 21이고 나머지 0이므로 합은 21.
위 시나리오에서 얻은 출력값 25, 26, 0, 21이 예제 출력과 일치함을 확인할 수 있다. 이로써 서브트리 및 경로에 대한 가산/곱 연산과 합 쿼리가 제대로 수행됨을 검증했다.