문제 65
트리 경로 최댓값
$1$번부터 $N$번까지 번호가 붙은 정점으로 이루어진 트리가 있다. 각 정점 $i$에는 정수 $A_i$가 적혀 있다.
$Q$개의 쿼리를 순서대로 처리하라. 쿼리는 다음 중 하나이다.
1 u x: $A_u$를 $x$로 바꾼다. $(1 \le u \le N; -10^9 \le x \le 10^9)$2 u v: $u$번 정점에서 $v$번 정점으로 가는 단순 경로에 포함된 정점에 적힌 값의 최댓값을 출력한다. $(1 \le u, v \le N)$
입력
첫 번째 줄에 두 정수 $N$, $Q$가 공백으로 구분되어 주어진다. $(1 \le N, Q \le 200\,000)$
두 번째 줄에 $N$개의 정수 $A_1, A_2, \cdots, A_N$이 공백으로 구분되어 주어진다. $(-10^9 \le A_i \le 10^9)$
세 번째 줄부터 $N-1$개의 줄에 걸쳐 간선으로 연결된 두 정점의 번호를 의미하는 두 정수 $u$, $v$가 공백으로 구분되어 주어진다. $(1 \le u, v \le N; u \neq v)$
그다음 줄부터 $Q$개의 줄에 걸쳐 쿼리가 하나씩 주어진다.
출력
2 u v 쿼리마다 경로 최댓값을 한 줄에 하나씩 출력한다.
예제 입력 1
5 7 3 -2 5 1 4 1 2 1 3 2 4 2 5 2 4 3 1 2 10 2 4 3 2 5 4 1 3 -5 2 3 5 2 1 1
예제 출력 1
5 10 10 10 3
공식 해설
트리의 경로를 배열 구간으로 나누기
정점의 값이 바뀌므로 경로의 최댓값을 미리 고정해 둘 수 없다. Heavy-Light Decomposition(HLD)으로 경로를 적은 수의 배열 구간으로 바꾸고 세그먼트 트리를 사용한다.
트리의 루트를 하나 정하고 각 정점의 서브트리 크기를 구한다. 자식 중 서브트리가 가장 큰 하나와의 간선을 무거운 간선으로 정하고, 나머지를 가벼운 간선으로 정한다. 무거운 자식을 먼저 방문하는 DFS 순서로 번호를 붙이면 무거운 간선으로 이어진 각 체인이 배열의 연속 구간이 된다.
가벼운 간선으로 부모에서 자식으로 내려가면 서브트리 크기가 절반 미만으로 줄어든다. 따라서 루트로 향하는 경로에서 체인이 바뀌는 횟수는 $O(\log(N+1))$이다.
갱신과 경로 질의
DFS 번호 위치에 해당 정점의 값을 넣고 구간 최댓값 세그먼트 트리를 만든다. 정점 값의 변경은 그 위치의 점 갱신이다.
$u,v$의 체인이 다르면 체인의 맨 위 정점이 더 깊은 쪽을 고른다. 그 체인의 맨 위부터 현재 정점까지의 최댓값을 구하고, 현재 정점을 체인 맨 위의 부모로 올린다. 두 정점이 같은 체인에 들어오면 둘 사이의 배열 구간을 마지막으로 조회한다.
이렇게 나눈 구간들은 원래 경로를 정확히 덮는다. 양 끝 정점도 포함하며 $u=v$이면 정점 하나의 값이다. 음수 값도 있으므로 답의 초기값은 음의 무한대로 둔다.
전처리는 $O(N)$, 점 갱신은 $O(\log(N+1))$, 경로 질의는 $O(\log^2(N+1))$이다. 전체 시간복잡도는 $O(N+Q\log^2(N+1))$, 공간복잡도는 $O(N)$이다.
구현 · 언어별 풀이
C++ 구현
부모·깊이·서브트리 크기·무거운 자식·체인 머리·DFS 번호를 vector<int>로 관리한다. 원래 정점 값을 DFS 번호에 배치한 뒤 세그먼트 트리를 만든다. 긴 경로 모양의 트리에서는 명시적 스택과 역순 순회로 서브트리 크기를 구하면 깊은 재귀 호출을 피할 수 있다.
C 구현
부모·깊이·크기·heavy·head·pos를 int 배열로 둔다. 반복 DFS 순서를 뒤집어 subtree size를 계산하고 heavy 경로를 먼저 배치한다. 깊이 200000의 재귀는 피하며 음수 값용 max 초기값을 둔다.
Python 구현
트리 전처리는 명시적 스택과 역순 순회로 수행한다. HLD 인덱스 배열과 구간 최댓값 트리를 두고 경로를 chain별로 나눈다. max의 기본값을 0으로 두면 음수 경로에서 틀린다.
Java 구현
int[] 부모·깊이·heavy·head·pos와 값용 트리를 사용한다. 전처리 DFS를 반복문으로 하여 긴 트리의 StackOverflowError를 피한다. 경로의 모든 값이 음수일 때도 실제 최댓값을 반환한다.
Rust 구현
전처리 순서를 Vec 스택으로 얻고 역순에서 크기와 heavy 자식을 정한다. head와 pos는 usize, 최댓값은 부호 있는 자료형으로 둔다. chain을 넘을 때 더 깊은 head 쪽부터 처리한다.
JavaScript 구현
부모·깊이·서브트리 크기·heavy 자식·체인 머리·순회 번호를 정수 배열로 관리한다. 길이 $200\,000$인 경로도 가능하므로 원래 트리 DFS는 명시적 스택으로 순서를 만들고, 역순으로 크기와 heavy 자식을 계산한다. 체인 번호를 붙이는 순회도 heavy 경로를 먼저 내려가는 반복문으로 처리한다.
원소와 최댓값은 $[-10^9,10^9]$ 안이므로 Int32Array에 들어간다. 비어 있는 세그먼트 트리 칸에는 모든 입력보다 작은 -2147483648을 넣는다. $0$을 기본 최댓값으로 사용하면 음수만 있는 경로의 답이 틀린다.
점 갱신은 pos[v]에 반영하고, 경로 조회에서는 깊은 쪽 체인 머리까지의 닫힌 구간을 처리한 뒤 그 머리의 부모로 올라간다. 같은 체인에서는 두 순회 번호 사이를 마지막으로 조회한다. 두 정점이 같아도 그 정점 값을 포함한다.