선분 트리 병합을 이용한 [FJOI2018] 지도자 그룹 문제 해결

먼저 이 문제는 그리디로 풀기 어려우므로 동적 계획법(DP)을 고려할 수 있습니다.

처음에는 \(dp_{i, j}\)를 \(i\)번 노드를 루트로 할 때 선택된 노드 중 최소 가중치가 \(j\)인 최대 멤버 수로 정의하는 직관적인 방법이 있습니다. 하지만 이 방법은 \(O(n^3)\)의 시간 복잡도를 가집니다. 많은 전이 과정이 동일하므로 비효율적입니다. 더 빠른 전이를 위해 상태를 변경하여, \(dp_{i, j}\)를 \(i\)를 루트로 하는 서브트리 내에서 선택된 모든 노드의 가중치가 \(j\) 이상일 때의 최대 멤버 수로 정의할 수 있습니다. 이 DP 값은 뒤에서 앞으로 갈수록 단조 증가합니다. 따라서 각 서브트리의 전이는 \(dp_{u, i} = \sum dp_{v, i}\)로 간단히 표현됩니다. 만약 이 집합에 노드 \(u\)를 포함한다면, \(dp_{u, i} = \max\{dp_{u, i}, dp_{u, w_u} + 1\}(i \le w_u)\)라는 추가 전이가 가능하며, 이를 통해 \(O(n^2)\)의 복잡도를 달성할 수 있습니다.

이제 이 DP를 최적화할 방법을 생각해보겠습니다. 첫 번째 전이 방정식은 서브트리의 각 위치 값을 단순히 더하는 것과 같습니다. 이는 선분 트리 병합의 과정과 매우 유사하므로, 선분 트리 병합의 관점에서 접근할 수 있습니다. 두 번째 전이 방정식의 경우, \(dp_i\)의 각 값 \(dp_{i, j}\)가 뒤에서 앞으로 갈수록 단조 증가하므로, 현재 노드의 선분 트리에서 특정 구간에 1을 더하는 것과 같습니다. 이는 영구적 지연 기법(lazy propagation)을 사용하여 해결할 수 있지만, 구간에 동일한 값을 더하고 최종적으로 한 번만 값을 조회하면 되므로, DP 값을 뒤에서 앞으로 차분(difference)하는 방법을 고려할 수 있습니다. 이 경우 \(w_i\) 위치에 1을 더하고, \(i\)의 왼쪽에서 첫 번째로 0이 아닌 위치에서 1을 빼면 됩니다. 이 위치는 선분 트리 이분 탐색을 통해 찾을 수 있습니다. 효율적인 구현 방법은 \(1 \sim w_i\) 구간의 합 \(S\)를 조회한 후, 선분 트리에서 이분 탐색을 통해 \(1 \sim w_i\) 구간의 합이 \(S\)가 되는 첫 번째 위치를 찾는 것입니다.

주의할 점:

  • 동적 할당 시 변수명(예: \(num\), \(tot\), \(cnt\))을 혼동하지 않도록 주의해야 합니다.

  • 선분 트리 병합 함수를 호출할 때는 노드 번호가 아닌 루트(\(rt\))를 전달해야 합니다.

#include<bits/stdc++.h>
using namespace std;
#define N 200000 + 5
#define M 4000000 + 5
#define ls t[p].l
#define rs t[p].r
#define mid (l + r >> 1)
#define rep(i, l, r) for(int i = l; i <= r; ++i)
#define Next(i, u) for(int i = h[u]; i; i = e[i].next)
struct edge{
    int v, next;
}e[N << 1];
struct tree{
    int l, r, sum;
}t[M];
int n, u, ok, tot, cnt, num, d[N], w[N], h[N], rt[N];
int read(){
    char c; int x = 0, f = 1;
    c = getchar();
    while(c > '9' || c < '0'){ if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
void add(int u, int v){
    e[++tot].v = v, e[tot].next = h[u], h[u] = tot;
    e[++tot].v = u, e[tot].next = h[v], h[v] = tot;
}
void update(int &p, int l, int r, int x, int y, int k){
    if(!p) p = ++num;
    if(l >= x && r <= y){ t[p].sum += k; return;}
    if(mid >= x) update(ls, l, mid, x, y, k);
    if(mid < y) update(rs, mid + 1, r, x, y, k);
    t[p].sum = t[ls].sum + t[rs].sum;
}
int query(int p, int l, int r, int x, int y){
    if(!p) return 0;
    if(l >= x && r <= y) return t[p].sum;
    int ans = 0;
    if(mid >= x) ans += query(ls, l, mid, x, y);
    if(mid < y) ans += query(rs, mid + 1, r, x, y);
    return ans;
}
int find(int p, int l, int r, int k){
    if(l == r) return !t[p].sum ? 0 : l;
    if(t[ls].sum >= k) return find(ls, l, mid, k);
    else return find(rs, mid + 1, r, k - t[ls].sum);
}
void Merge(int &p, int k, int l, int r){
    if(!p || !k){ p = p + k; return;}
    if(l == r){ t[p].sum += t[k].sum; return;}
    Merge(ls, t[k].l, l, mid), Merge(rs, t[k].r, mid + 1, r);
    t[p].sum = t[ls].sum + t[rs].sum;
}
void dfs(int u, int fa){
    Next(i, u){
        int v = e[i].v; if(v == fa) continue;
        dfs(v, u), Merge(rt[u], rt[v], 1, cnt);
    }
    int val = query(rt[u], 1, cnt, 1, w[u]);
    int pos = find(rt[u], 1, cnt, val);
    update(rt[u], 1, cnt, w[u], w[u], 1);
    if(pos) update(rt[u], 1, cnt, pos, pos, -1);
}
int main(){
    n = read();
    rep(i, 1, n) w[i] = d[i] = read();
    sort(d + 1, d + n + 1);
    cnt = unique(d + 1, d + n + 1) - d - 1;
    rep(i, 1, n) w[i] = lower_bound(d + 1, d + cnt + 1, w[i]) - d;
    rep(i, 2, n) u = read(), add(u, i);
    dfs(1, 0);
    printf("%d", t[rt[1]].sum);
    return 0;
}

이 차분 배열을 유지하는 또 다른 방법은 균형 이진 트리 휴리스틱 병합(heuristic merge)을 사용하는 것입니다. C++의 set을 활용하면 구현이 매우 간단해집니다. 자식 노드의 차분 배열을 병합할 때는 휴리스틱 병합을 사용하여 무식하게 삽입합니다. 단일 위치 변경을 위해 set에 각 원소 \(x\)를 저장하여 위치 \(x\)에서 차분 배열이 +1임을 나타냅니다. \(w_i\)를 삽입할 때마다 \(w_i\)보다 작은 첫 번째 원소를 찾아 제거하면 됩니다. 최종 답은 1번 노드의 set 크기입니다.

#include<bits/stdc++.h>
using namespace std;
#define N 200000 + 5
#define rep(i, l, r) for(int i = l; i <= r; ++i)
#define Next(i, u) for(int i = h[u]; i; i = e[i].next)
struct edge{
    int v, next;
}e[N << 1];
multiset <int> S[N];
multiset <int> :: iterator it;
int n, u, tot, cnt, h[N], d[N], w[N];
int read(){
    char c; int x = 0, f = 1;
    c = getchar();
    while(c > '9' || c < '0'){ if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
void add(int u, int v){
    e[++tot].v = v, e[tot].next = h[u], h[u] = tot;
    e[++tot].v = u, e[tot].next = h[v], h[v] = tot;
}
void Merge(int x, int y){
    if(S[x].size() < S[y].size()) swap(S[x], S[y]);
    for(it = S[y].begin(); it != S[y].end(); ++it) S[x].insert(*it);
}
void dfs(int u, int fa){
    Next(i, u){
        int v = e[i].v; if(v == fa) continue;
        dfs(v, u), Merge(u, v);
    }
    S[u].insert(w[u]);
    it = S[u].lower_bound(w[u]); 
    if(it != S[u].end() && it != S[u].begin()) S[u].erase(--it);
}
int main(){
    n = read();
    rep(i, 1, n) w[i] = d[i] = read();
    sort(d + 1, d + n + 1);
    cnt = unique(d + 1, d + n + 1) - d - 1;
    rep(i, 1, n) w[i] = lower_bound(d + 1, d + cnt + 1, w[i]) - d;
    rep(i, 2, n) u = read(), add(u, i);
    dfs(1, 0);
    printf("%d", S[1].size());
    return 0;
}

이와 같은 문제에는 일반적인 패턴이 있습니다:

DP 값이 일차 함수나 상수 함수의 형태를 띨 때, 먼저 차분을 한 후 set이나 선분 트리 병합을 사용하여 유지할 수 있습니다. 선분 트리는 구간 덧셈을 지원하므로, 구간 수정이 필요한 경우 선분 트리 병합을 사용합니다. 반면, AT2347 [ARC070C] NarrowRectangles와 같은 문제는 set을 사용하여 변곡점을 유지하고 함수의 평행 이동을 지원할 수 있습니다.

태그: DP 선분트리 세그먼트트리병합 차분배열 휴리스틱병합

6월 25일 01:53에 게시됨