Heavy-Light Decomposition Recursive DP
algorithm , tree , dynamic-programming , heavy-light-decomposition
개요
Heavy-Light Decomposition
Heavy-Light Decomposition(HLD)은 루트 있는 트리를 여러 개의 heavy chain으로 분할하는 알고리즘으로, 다음의 성질들을 만족합니다.
-
각 간선은 heavy edge 또는 light edge로 분류됩니다.
-
리프가 아닌 정점은 자식들 중 정확히 하나를 heavy child로 갖습니다. heavy child로 가는 간선은 heavy edge이고, 다른 자식으로 가는 간선은 모두 light edge입니다.
-
heavy edge로 연결된 정점들의 묶음을 heavy chain이라 합니다. 각 heavy chain의 형태는 한 정점에서 출발해서 리프까지 내려가는 연속적인 경로가 되고, 모든 정점은 정확히 하나의 heavy chain에 속하게 됩니다.
-
트리 위의 어떤 경로에 대해서도 경로 상의 light edge의 개수가 $O(\log N)$개입니다. 여기에서 $N$은 정점의 개수입니다.
HLD는 잘 알려진 알고리즘이고 이미 다른 글에서도 많이 다루고 있기 때문에 추가적인 설명이나 구현 방법은 생략하겠습니다. HLD에 대한 더 자세한 설명은 다음 글에서 읽어 보실 수 있습니다.
HLD는 주로 트리에서의 경로 쿼리를 처리하기 위해 사용되고, 삼성소프트웨어멤버십 블로그에도 이와 관련된 테크닉을 소개하는 글이 존재합니다. 하지만 HLD는 경로 쿼리 외에도 다양한 활용처가 존재합니다.
Heavy-Light Decomposition Recursive DP
이 글에서는 HLD를 사용하는 새로운 테크닉인 Heavy-Light Decomposition Recursive DP (HLRecDP)를 소개합니다. 비교적으로 최근에 2018년 논문으로 소개되었고 프로그래밍 대회에도 몇 번 등장했지만, 여러 문제에 활용될 가능성과 구현의 간결함에 비해 거의 알려져 있지 않고 한국어 자료도 없어서 글을 작성하게 되었습니다.
HLRecDP는 트리 위에서 특정한 제약이 붙은 knapsack DP의 시간복잡도를 최적화하는 데에 사용됩니다. HLRecDP를 사용한 풀이는 $O(N^{1.58}X)$처럼 독특한 시간복잡도를 갖는 경우가 많습니다.
아래의 트리 독립 집합 냅색 문제를 예시로 들어서 살펴보겠습니다.
트리 독립 집합 냅색 문제
문제
정점이 $N$개인 트리가 주어지고, 배낭의 용량 $X$가 주어집니다. $(1 \leq N \leq 200, 1 \leq X \leq 50,000)$
정점 $1 \leq i \leq N$는 무게 $w_i$와 가치 $v_i$라는 값을 가지고 있습니다. $(1 \leq w_i \leq X, 1 \leq v_i \leq 10^7)$
각 정점을 배낭에 넣을지 말지를 정해서, 무게 합이 배낭의 용량 이하가 되도록 하면서 가치 합을 최대화하고 싶습니다. 이 때, 배낭에 넣은 정점들끼리는 간선으로 인접해 있으면 안 됩니다.
다시 말해, 트리 $T = (V,E)$에 대해 정점 부분집합 $S \subseteq V$를 결정해서, 아래의 두 제약을 모두 만족하는 $\sum_{i \in S} v_i$의 최댓값을 구해야 합니다.
-
$\sum_{i \in S} w_i \leq X$ (무게 상한 제약)
-
$a,b \in S$라면 $(a,b) \not \in E$ (독립 집합 제약)
이 글의 목표는 트리 독립 집합 냅색 문제에 대해 2초 이내에 동작하는 코드를 작성하는 것입니다.
풀이 1
트리 DP
루트를 1번 정점으로 고정한 뒤 간단한 트리 DP를 통해 $O(NX^2)$에 해결할 수 있습니다.
각 정점 $n$마다 다음과 같은 DP를 정의합니다.
$D[n][i]$ : 정점 $n$의 서브트리에서 무게 상한 제약이 $i$일 때, 트리 독립 집합 냅색 문제의 답
$E[n][i]$ : 정점 $n$의 서브트리에서 무게 상한 제약이 $i$일 때, $n$을 배낭에 넣지 않는 트리 독립 집합 냅색 문제의 답
정점 $n$의 자식들이 $n_1, n_2, \cdots, n_k$일 때, 상태 전이는 아래와 같습니다.
$D[n][i] = \max(\max_{i_1 + i_2 + \cdots + i_k = i} (\sum_{1 \leq j \leq k} D[n_j][i_j]), \max_{i_1 + i_2 + \cdots + i_k = i-w_n} (\sum_{1 \leq j \leq k} E[n_j][i_j]) + v_n)$
$E[n][i] = \max_{i_1 + i_2 + \cdots + i_k = i} (\sum_{1 \leq j \leq k} D[n_j][i_j])$
원래 문제의 정답은 $D[1][X]$가 됩니다.
max-plus convolution
편의를 위해 max-plus convolution이라는 새로운 연산을 정의합시다.
$\textrm{max-plus-convolution}(A, B)$는 길이가 $X+1$인 두 배열 $A,B$를 입력으로 받아 길이가 $X+1$인 배열 $C$를 반환하는 함수입니다. 이 때 $C_i = \max_{0 \leq j \leq i} (A_j + B_{i-j})$로 정의됩니다.
max-plus convolution의 시간복잡도는 $O(X^2)$이고 위 DP의 상태 전이는 max-plus convolution를 $O(N)$번 하는 형태이므로, 최종 시간복잡도는 $O(NX^2)$가 됩니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
int N, X;
int W[303], V[303];
vector<int> g[303];
vector<int> max_plus_convolution(const vector<int> &A, const vector<int> &B)
{
vector<int> C(X+1);
for(int i=0; i<=X; i++)
for(int j=0; j<=i; j++)
C[i] = max(C[i], A[j] + B[i-j]);
return C;
}
pair<vector<int>, vector<int> > DFS(int n, int prev)
{
vector<int> D(X+1), E(X+1);
for(int next : g[n])
{
if (next == prev) continue;
auto [DD,EE] = DFS(next, n);
D = max_plus_convolution(D, DD);
E = max_plus_convolution(E, EE);
}
for(int i=X; i>=0; i--)
{
E[i] = D[i];
if (i >= W[n]) D[i] = max(D[i], E[i-W[n]] + V[n]);
}
return {D, E};
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> X;
for(int i=1; i<=N; i++)
cin >> W[i] >> V[i];
for(int i=0; i<N-1; i++)
{
int a, b;
cin >> a >> b;
g[a].push_back(b);
g[b].push_back(a);
}
auto [D,E] = DFS(1, 0);
cout << D[X] << "\n";
return 0;
}
풀이 2
위의 DP 풀이는 $X$ 제한이 크다면 너무 느리게 작동한다는 문제점이 있습니다.
만약 독립 집합 제약이 없다면 평범한 냅색 문제와 동일해지고, 트리 대신에 배열로 생각해서 $O(NX)$에 해결할 수 있습니다.
냅색 문제의 풀이를 트리에 그대로 적용할 수 없는 이유는, 한 정점이 둘 이상의 자식을 가질 때 두 자식에 대한 DP 배열을 합쳐 주는 max-plus convolution 연산이 $O(X^2)$으로 느리기 때문입니다. max-plus convolution을 최적화해서 시간복잡도를 개선할 수 있을까요?
아쉽게도 일반적인 (convex 같은 제약이 없는) max-plus convolution은 매우 어려운 문제라서 $O(X^2)$보다 빠른 알고리즘은 복잡하면서도 competitive programming에서 비실용적입니다. 따라서 max-plus convolution을 사용하지 않는 새로운 풀이를 생각해야 합니다.
Recursive DP
여기에서 새로운 테크닉인 Recursive DP를 소개합니다. 아이디어는 트리를 수많은 직선들로 펼쳐서 max-plus convolution을 사용하지 않고 냅색 문제의 풀이와 동일한 상태 전이만으로 백트래킹하듯이 문제를 해결하는 것입니다.
이를 위해, 앞으로 이어서 계산하고 싶은 DP 배열의 초깃값을 DFS 함수의 인자로 넘겨주고, 넘겨받은 DP 배열에 새로운 값을 누적해서 반환하는 방식을 사용합니다. 일반적으로 부르는 “재귀 DP”와는 다른 테크닉으로, 상당히 비직관적으로 느껴질 수 있으니 주의 깊게 보시길 바랍니다.
코드
DP의 정의가 달라졌다는 사실에 유의해야 합니다. 원래의 $D[n]$은 정점 $n$의 서브트리에 대한 정답을 담고 있었지만, Recursive DP의 $D[n]$은 $n$을 방문하기 이전에 DFS에서 순회가 끝난 모든 정점(=조상의 왼쪽 자손들)에 대한 정답에 추가해서, 정점 $n$의 서브트리에 대한 정답을 누적한 형태의 값을 담고 있습니다.
자식으로 내려갈 때마다 $D[n]$을 구하기 위해, $E[n]$을 구하기 위해 총 두 번의 재귀 호출을 하므로 재귀 함수가 트리의 깊이에 대해 지수적으로 늘어나게 되고, 시간복잡도는 $O(2^N X)$가 됩니다.
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
int N, X;
int W[303], V[303];
vector<int> g[303];
pair<vector<int>, vector<int> > DFS(int n, int prev, const vector<int> &dp)
{
auto D = dp, E = dp;
for(int next : g[n])
{
if (next == prev) continue;
D = DFS(next, n, D).first;
E = DFS(next, n, E).second;
}
for(int i=X; i>=0; i--)
{
E[i] = D[i];
if (i >= W[n]) D[i] = max(D[i], E[i-W[n]] + V[n]);
}
return {D, E};
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> X;
for(int i=1; i<=N; i++)
cin >> W[i] >> V[i];
for(int i=0; i<N-1; i++)
{
int a, b;
cin >> a >> b;
g[a].push_back(b);
g[b].push_back(a);
}
auto [D,E] = DFS(1, 0, vector<int>(X+1));
cout << D[X] << "\n";
return 0;
}
풀이 3
균형 트리의 Recursive DP
모든 정점에 대해 자식의 서브트리 크기가 항상 자신의 서브트리 크기의 절반 이하인 경우 균형 트리라고 정의합시다.
균형 트리에 대해서는 위 Recursive DP 풀이의 시간복잡도가 $O(2^N X)$보다 작아집니다.
$f(n)$을 정점 개수가 $n$개인 서브트리의 루트 $r$의 $D[r],E[r]$을 계산하는 시간복잡도라 하면 아래의 수식으로 나타낼 수 있습니다. DFS의 인자로 DP 배열을 넘겨줄 때 const reference로 넘겨주고 std::move
로 돌려받으면 $O(1)$에 주고받을 수 있음을 참고하면 좋습니다.
$f(n) \leq 2(f(n_1) + f(n_2) + \cdots + f(n_k)) + O(X), n_1 + n_2 + \cdots + n_k = n-1$
$f(n)$은 $n$에 대한 1차 이상의 다항식이므로 볼록하기 때문에 위 수식은 $n_1 = \lceil \frac{n-1}{2} \rceil, n_2 = \lfloor \frac{n-1}{2} \rfloor$일 때 최대화되고, 다시 아래처럼 풀어 적을 수 있습니다.
$f(n) \leq 4f(\frac{n}{2}) + O(X)$
따라서 $f(N) = O(N^2 X)$가 됩니다.
HLD + Recursive DP
위의 시간복잡도 분석은 그대로 일반적인 트리에 적용할 수 없지만, 위 아이디어를 응용해서 일반적인 트리에서도 시간복잡도를 개선할 수 있습니다.
HLD를 사용하면 경로 상의 light edge의 개수가 $O(\log N)$개로 작음을 이용해서, heavy child로 내려갈 때에는 가벼운 비용을 지불하고 light child로 내려갈 때에만 무거운 비용을 지불하게 설계할 것입니다.
핵심 관찰: Recursive DP의 코드에서 첫 번째 자식으로의 두 번의 재귀 호출은 초기 상태에서 $D = E$이므로 동일한 DP 배열을 인자로 넘겨주게 됩니다. 따라서 한 번의 재귀 호출만으로 충분합니다.
HLD를 사용한 뒤 첫 번째 자식을 heavy child로 고정하고 다시 시간복잡도 분석을 해 봅시다.
$f(n) \leq f(n_1) + 2(f(n_2) + \cdots + f(n_k)) + O(X), n_1 + n_2 + \cdots + n_k = n-1, n_1 \geq n_2 \geq \cdots \geq n_k$
위 수식은 $n_1 = \lceil \frac{n-1}{2} \rceil, n_2 = \lfloor \frac{n-1}{2} \rfloor$일 때 최대화되고, 다시 아래처럼 풀어 적을 수 있습니다.
$f(n) \leq 3f(\frac{n}{2}) + O(X)$
따라서 $f(N) = O(N^{\log_2 3} X) \approx O(N^{1.58} X)$가 됩니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
int N, X;
int W[303], V[303], sz[303];
vector<int> g[303];
void HLD(int n, int prev)
{
g[n].erase(remove(g[n].begin(), g[n].end(), prev), g[n].end());
sz[n] = 1;
for(int next : g[n])
{
HLD(next, n);
sz[n] += sz[next];
}
sort(g[n].begin(), g[n].end(), [&](int a, int b){
return sz[a] > sz[b];
});
}
pair<vector<int>, vector<int> > DFS(int n, const vector<int> &dp)
{
auto D = dp, E = dp;
if (!g[n].empty())
{
tie(D, E) = DFS(g[n][0], dp);
for(int i=1; i<g[n].size(); i++)
{
D = DFS(g[n][i], D).first;
E = DFS(g[n][i], E).second;
}
}
for(int i=X; i>=0; i--)
{
E[i] = D[i];
if (i >= W[n]) D[i] = max(D[i], E[i-W[n]] + V[n]);
}
return {D, E};
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> X;
for(int i=1; i<=N; i++)
cin >> W[i] >> V[i];
for(int i=0; i<N-1; i++)
{
int a, b;
cin >> a >> b;
g[a].push_back(b);
g[b].push_back(a);
}
HLD(1, 0);
auto [D,E] = DFS(1, vector<int>(X+1));
cout << D[X] << "\n";
return 0;
}
연습 문제
Protect the Pollen!
vertex cover의 여집합이 independent set과 동치이므로, 트리 독립 집합 냅색 문제와 완전히 동일한 문제가 됩니다.
제한이 작아서 $O(NS^2)$ 풀이가 통과하지만, 만약 $S$ 제한이 $10^5$로 늘어나더라도 HLRecDP로 $O(N^{1.58}S)$에 해결할 수 있습니다.
코드는 위와 같으므로 생략합니다.
2018 Japan Domestic Contest H. For Programming Excellence
문제
정점이 $N$개이고 루트가 1번 정점인 트리가 주어지고, 분배 가능한 레벨 총합의 상한 $K$가 주어집니다. $(2 \leq N \leq 100, 1 \leq K \leq 10^5)$
각 정점 $1 \leq i \leq N$은 레벨 상한 $1 \leq h_i \leq 10^5$, 중요도 $1 \leq s_i \leq 10^9$, 부모 정점 $1 \leq p_i < i$, 부모의 레벨 하한 제약 $1 \leq l_i \leq h_{p_i}$라는 값들을 가지고 있습니다. $p_1, l_1$ 값은 존재하지 않습니다.
각 정점 $i$의 레벨을 $x_i$라 하면, 처음에는 모든 $i$에 대해 $x_i = 0$입니다.
$x_i$들을 잘 결정해서 아래 제약들을 모두 만족하는 $\sum_{1 \leq i \leq N} x_i s_i$의 최댓값을 구해야 합니다.
-
$\sum_{1 \leq i \leq N} x_i \leq K$ (레벨 총합의 상한 제약)
-
$1 \leq i \leq N$에서 $0 \leq x_i \leq h_i$ (레벨 상한 제약)
-
$2 \leq i \leq N, x_i > 0$에서 $x_{p_i} \geq l_i$ (부모의 레벨 하한 제약)
풀이
다음과 같은 DP를 정의합니다.
$dp[n][i]$ : 정점 $n$의 서브트리에서 레벨 총합의 상한 제약이 $i$이고 $n$번 정점의 부모의 레벨 하한 제약이 만족된 상태일 때, (레벨) $\times$ (중요도) 합의 최댓값
$n$의 자식들을 $l_i$의 오름차순으로 정렬해 두면, $x_n$ 값을 고정했을 때 부모의 레벨 하한 제약이 만족되는 자식들은 prefix로만 나타납니다. 따라서 가능한 모든 prefix에 대해, 이 prefix의 정점들만 부모의 레벨 하한 제약이 만족되기 위해 요구되는 $x_n$의 구간 $[L,R]$(query interval)을 전처리할 수 있습니다.
상태 전이를 prefix마다 독립적으로 생각해 보겠습니다. prefix의 자식들이 $n_1, n_2, \cdots, n_k$이고 query interval이 $[L,R]$이라 했을 때의 점화식은 아래와 같습니다.
$dp[n][i] = \max_{i-R \leq j \leq i-L} (\max_{i_1 + i_2 + \cdots + i_k = j}(\sum_{1 \leq t \leq k} (dp[n_t][i_t])) + s_n \cdot (i-j))$
자식들을 앞에서부터 차례대로 방문하는 recursive DP를 작성하면 안쪽의 max-plus convolution을 지울 수 있습니다. 바깥쪽의 range max query는 $i$의 오름차순으로 계산한다면 deque DP라는 잘 알려진 테크닉으로 최적화할 수 있습니다. 따라서 $dp[n][i]$를 $O(1)$에 구할 수 있고, 최종 시간복잡도는 $O(NK)$가 됩니다.
이 문제에서는 구해야 하는 DP 배열이 두 개가 아니라 하나뿐이기 때문에 자식마다 한 번의 재귀 호출만을 하게 되고, 시간복잡도가 지수적으로 늘어나지 않아 HLD를 사용할 필요가 없습니다. 심지어 DFS 함수의 인자인 DP 배열이 항상 하나의 근원만을 갖기 때문에, DP 배열을 reference로 넘겨줄 필요가 없이 전역 배열로 관리해도 충분합니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
// deque로 range max query 최적화
struct Deque {
ll X[100001];
int T[100001];
int l = 0, r = 0; // [l, r)
void insert(ll x, int t) {
while(l < r && X[r-1] < x) r--;
X[r] = x, T[r] = t;
r++;
}
void erase(int t) {
while(l < r && T[l] < t) l++;
}
ll query() {
return X[l];
}
void clear() {
l = r = 0;
}
} dq;
int N, K;
int h[101], s[101], l[101];
vector<int> g[101];
vector<ll> dp;
void DFS(int n)
{
sort(g[n].begin(), g[n].end(), [&](int a, int b){
return l[a] < l[b];
});
// query interval 구하기
list<pair<int, int> > v(1, {0, h[n]});
for(int next : g[n])
{
v.back().second = l[next]-1;
v.push_back({l[next], h[n]});
}
// ndp[i] = chmax_{i-R <= j <= i-L}(dp[j] + s[n]*(i-j))
vector<ll> ndp(K+1, 0);
auto UpdateDP = [&]() {
auto [L,R] = v.front(); v.pop_front();
if (L > R) return;
dq.clear();
for(int i=L,work=0; i<=K; i++)
{
for(int &j=work; j<=i-L; j++)
dq.insert(dp[j] - (ll)s[n]*j, j);
dq.erase(i-R);
ndp[i] = max(ndp[i], dq.query() + (ll)s[n]*i);
}
};
UpdateDP();
for(int next : g[n])
{
DFS(next);
UpdateDP();
}
dp = move(ndp);
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
while(1)
{
cin >> N >> K;
if (N == 0 && K == 0) break;
// clear
for(int i=1; i<=N; i++)
g[i].clear();
for(int i=1; i<=N; i++)
cin >> h[i];
for(int i=1; i<=N; i++)
cin >> s[i];
for(int i=2,p; i<=N; i++)
{
cin >> p;
g[p].push_back(i);
}
for(int i=2; i<=N; i++)
cin >> l[i];
dp = vector<ll>(K+1, 0);
DFS(1);
cout << dp[K] << "\n";
}
return 0;
}
AtCoder Beginner Contest 311 Ex. Many Illumination Plans
문제에서 주어진 연산을 간단하게 변형한 형태로 서술하겠습니다.
문제
정점이 $N$개이고 루트가 1번 정점인 트리가 주어지고, 배낭의 용량 $X$가 주어집니다. $(2 \leq N \leq 200, 0 \leq X \leq 50,000)$
정점 $1 \leq i \leq N$는 가치 $B_i$, 무게 $W_i$, 색깔 $C_i$를 가지고 있습니다. $(0 \leq B_i \leq 10^{15}, 0 \leq W_i \leq X, 0 \leq C_i \leq 1)$
다음 문제에 대한 정답을 $R(v)$라 정의하겠습니다.
루트가 $v$인 서브트리에 대해서만 고려하겠습니다. 각 정점을 배낭에 넣을지 말지를 정해서, 무게 합이 배낭의 용량 이하가 되도록 하면서 가치 합을 최대화하고 싶습니다. 이 때 루트인 $v$는 반드시 배낭에 넣어야 하고, 배낭에 넣은 정점들만으로 이루어진 압축된 트리에서 인접한 정점끼리는 색깔이 달라야 합니다.
$R(1), R(2), \cdots, R(N)$을 모두 구해야 합니다.
풀이
특정 $v$에 대해 $R(v)$를 구하는 작업은 트리 독립 집합 냅색 문제와 거의 동일하게 DP 배열 $D[n][i], E[n][i]$를 정의하고 $O(N^{1.58}X)$에 해결할 수 있습니다. 하지만, 단순히 모든 $v$에 대해 $R(v)$를 독립적으로 구한다면 $O(N^{2.58}X)$으로 시간 초과를 받게 됩니다.
이를 개선하기 위해 HLD의 성질을 이용해서 모든 $v$에 대한 답을 한꺼번에 구할 것입니다. 핵심 아이디어는 heavy chain 단위로 묶어서 답을 계산하는 것입니다.
어떤 heavy chain의 루트에 HLRecDP를 호출하는 상황을 생각해 봅시다. heavy child로의 재귀 호출은 자신과 동일한 DP 배열을 인자로 넘겨 주게 되므로, 자신이 인자로 갖고 있었던 DP 배열이 초기 상태였다면 heavy child가 구한 정답을 활용해서 자신의 정답도 구할 수 있습니다. 반면에 light child가 구한 정답은 재활용할 수 없습니다.
정점이 $n$개인 서브트리 내부에 있는 모든 정점에 대한 답을 구하는 재귀 함수의 시간복잡도를 $g(n)$이라 정의하고, $f(n) = O(n^{1.58}X)$라 합시다. HLD를 사용한 뒤 첫 번째 자식을 heavy child로 고정하고 시간복잡도 분석을 해 봅시다.
$g(n) \leq g(n_1) + \cdots + g(n_k) + 2(f(n_2) + \cdots + f(n_k)) + O(X), n_1 + \cdots + n_k = n-1, n_1 \geq \cdots \geq n_k$
위 수식은 $n_1 = \lceil \frac{n-1}{2} \rceil, n_2 = \lfloor \frac{n-1}{2} \rfloor$일 때 최대화되고, 다시 아래처럼 풀어 적을 수 있습니다.
$g(n) \leq 2g(\frac{n}{2}) + 2f(\frac{n}{2}) + O(X)$
$g(n) \leq 2g(\frac{n}{2}) + O(n^{1.58}X)$
따라서 $g(N) = O(N^{1.58} X)$가 됩니다.
코드
코드는 아래와 같습니다.
solve_all = false
일 경우 일반적인 HLRecDP입니다.
solve_all = true
일 경우 서브트리 내부에 있는 모든 정점에 대한 정답을 구합니다. heavy child가 구한 정답에 해당하는 DP는 재활용해 주고, light child가 구한 정답에 해당하는 DP는 그냥 버리는 것을 확인할 수 있습니다. 하나의 heavy chain에 있는 정점들은 연쇄적으로 정보를 넘겨받으며 답을 구하게 됩니다.
#include<bits/stdc++.h>
using namespace std;
#define INF ((ll)1e18)
#define ll long long
int N, X;
vector<int> g[201];
ll B[201], res[201];
int p[201], sz[201], W[201], C[201];
pair<vector<ll>, vector<ll> > DFS(bool solve_all, int n, const vector<ll> &dp)
{
auto D = dp, E = dp;
if (!g[n].empty())
{
tie(D, E) = DFS(solve_all, g[n][0], dp);
for(int i=1; i<g[n].size(); i++)
{
D = DFS(false, g[n][i], D).first;
E = DFS(false, g[n][i], E).second;
if (solve_all) DFS(true, g[n][i], dp);
}
}
for(int i=W[n]; i<=X; i++)
{
if (C[n] == 0) D[i] = max(D[i], E[i-W[n]] + B[n]);
else E[i] = max(E[i], D[i-W[n]] + B[n]);
}
if (solve_all) res[n] = (C[n] == 0 ? (B[n] + E[X-W[n]]) : (B[n] + D[X-W[n]]));
return {D, E};
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> X;
for(int i=2; i<=N; i++)
{
cin >> p[i];
g[p[i]].push_back(i);
}
for(int i=1; i<=N; i++)
cin >> B[i] >> W[i] >> C[i];
for(int i=N; i>=1; i--) // HLD
{
sz[p[i]] += ++sz[i];
sort(g[i].begin(), g[i].end(), [&](int a, int b){
return sz[a] > sz[b];
});
}
DFS(true, 1, vector<ll>(X+1));
for(int i=1; i<=N; i++)
cout << res[i] << "\n";
return 0;
}