[알고리즘] 금광 세그먼트 트리(구간 최대 부분합)

금광 세그는 구간 [ L, R ] 에서 연속합의 최대를 구해주는 세그먼트 트리이다.

 

이름이 왜 금광이냐 하면

KOI 2014년도 중등부 4번 문제로 '금광'이란 문제가 나왔고 이때 만점자가 한 명도 없었다.

그때부터 이 테크닉이 유명해지면서 금광 세그라고 불리기 시작했다.

 

이 세그먼트 트리는 4가지의 값을 저장한다.

1. 왼쪽부터 시작하는 연속합의 최대
2. 오른쪽부터 시작하는 연속합의 최대
3. 그냥 구간합
4. 우리가 원하는 구간 연속합의 최대

 

1,2,3을 가지고 4를 구하는 느낌이라 보면 된다.

아래의 수열로 예시를 들어보자.

{10, -4, 3, 1, 5, 6, -35, 12, 21, -1}

1. 왼쪽부터 시작하는 연속합의 최대

왼쪽부터 시작하는 연속합은 {10, 6, 9, 10, 15, 21, -14, -2, 19, 18}이 있을 것이고

그 중 최대는 21이다.

2. 오른쪽부터 시작하는 연속합의 최대

오른쪽부터 시작하는 연속합은 {-1, 20, 32, -3, 3, 8, 9, 12, 8, 18}이 있을 것이고

그 중 최대는 32이다.

3. 그냥 구간합

sum({10, -4, 3, 1, 5, 6, -35, 12, 21, -1}) = 18

4. 우리가 원하는 구간 연속합의 최대

잘 찾아보면 12+21 = 33이다.

 

"구간 [ a, b ]와 [ b+1, c ] 를 합칠 때 1,2,3,4를 어떻게 갱신하는가" 도 알아야 한다.

[ a, b ]를 A, [ b+1, c ]를 B라고 하자.

합칠 때 가능한 경우들을 생각해서 그 중 최대를 골라 갱신할 것이다.

 

1. 왼쪽부터 시작하는 연속합의 최대

- A에서만 최대일 때

- A+B에서까지 최대일 때

new_left_max = max( A.left_max, A.sum + B.left_max );

2. 오른쪽부터 시작하는 연속합의 최대

- B에서만 최대일 때

- B+A에서까지 최대일 때

new_right_max = max( B.right_max, B.sum + A.right_max );

3. 그냥 구간합

이건 뭐 그냥 세그먼트 트리랑 똑같으니까 넘어가자

 

4. 우리가 원하는 구간 연속합의 최대

- A에서만 존재할 때

- B에서만 존재할 때

- A+B에 존재할 때

new_MAX = max({ A.MAX, B.MAX, A.right_max + B.left_max });

두 구간을 합치는 방법까지 알았으니 이제 코드로 넘어가 보자.

 

일단 아까 설명한 두 구간을 합치는 함수를 만들어 준다.

node merge(node A, node B){
    node NEW;
    NEW.left_max = max(A.left_max, A.sum + B.left_max);
    NEW.right_max = max(B.right_max, B.sum + A.right_max);
    NEW.MAX = max({A.MAX, B.MAX, A.right_max + B.left_max});
    NEW.sum = A.sum + B.sum;
    return NEW;
}

update와 query는 일반 세그먼트 트리와 거의 똑같다.

금광 세그먼트 트리를 보고 있는 시점에서 세그먼트 트리는 이미 알 것이기 때문에

update와 query는 설명하지 않겠다.

void update(int x, int i, int l, int r, int val){
    if(i<l || r<i) return;
    if(l == r){
        sg[x] = {val,val,val,val};
        return;
    }
    int m=(l+r)/2;
    update(x*2,i,l,m,val);
    update(x*2+1,i,m+1,r,val);
    sg[x] = merge(sg[x*2],sg[x*2+1]);
}
node query(int x, int s, int e, int l, int r){
    if(e < l || r < s) return {-INF,-INF,-INF,-INF};
    if(l<=s && e<=r) return sg[x];
    int m=(s+e)/2;
    node tmp1 = query(x*2,s,m,l,r);
    node tmp2 = query(x*2+1,m+1,e,l,r);
    return merge(tmp1,tmp2);
}

전체코드

#include<bits/stdc++.h>
using namespace std;
#define N 100100
typedef long long ll;
struct node{ll left_max,right_max,MAX,sum;} sg[N*4];
ll n,INF=1e15;
node merge(node A, node B){
    node NEW;
    NEW.left_max = max(A.left_max, A.sum + B.left_max);
    NEW.right_max = max(B.right_max, B.sum + A.right_max);
    NEW.MAX = max({A.MAX, B.MAX, A.right_max + B.left_max});
    NEW.sum = A.sum + B.sum;
    return NEW;
}
void update(int x, int i, int l, int r, int val){
    if(i<l || r<i) return;
    if(l == r){
        sg[x] = {val,val,val,val};
        return;
    }
    int m=(l+r)/2;
    update(x*2,i,l,m,val);
    update(x*2+1,i,m+1,r,val);
    sg[x] = merge(sg[x*2],sg[x*2+1]);
}
node query(int x, int s, int e, int l, int r){
    if(e < l || r < s) return {-INF,-INF,-INF,-INF};
    if(l<=s && e<=r) return sg[x];
    int m=(s+e)/2;
    node tmp1 = query(x*2,s,m,l,r);
    node tmp2 = query(x*2+1,m+1,e,l,r);
    return merge(tmp1,tmp2);
}
int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    n = 10;
    int arr[N] = {0,10,-4,3,1,5,6,-35,12,21,-1};
    for(int i=1; i<=n; i++)
        update(1,i,1,n,arr[i]);
    cout << query(1,1,n,1,10).MAX << '\n';
    return 0;
}

연습 문제

 

백준 16993 연속합과 쿼리

위의 코드에서 main만 조금 바꿔주면 AC 받을 수 있다.

 

백준 10167 금광

금광 세그라는 대명사가 된 문제이다.

힌트는 더보기에 있다.

좌표들을 x축이든 y축이든 정렬 후 무언가를 하는데 좌표가 겹칠 때가 있다.

이 좌표들을 한 번에 처리해 주는 게 포인트

시간 복잡도는 O(n^2 log n)이다.