구간 내 부분합 최대 세그먼트 트리 (금광세그)

문제

16993번: 연속합과 쿼리 (acmicpc.net)

16993번: 연속합과 쿼리

개요

어떻게 풀어야할까?

우선 누적합을 구하고 

그것을 이용해 최대값을 찾으려면 O(N^2)이 걸릴것이다. 시간초과가 날 복잡도이다.

이를 풀기위해선 구간 내 부분합 최대 세그먼트 트리(줄여서 금광세그) 자료구조를 사용해서 O(NlogN)으로 풀어야한다.

세그먼트 트리는 다음과 같은 자료구조로 이해할 수 있다.

F를 구간에 대해 알고싶은 값, 두 연속된 구간을 A,B라하면

F(A~B) = F(F(A),F(B)) 일때 쓸 수 있다.

예를들어 구간합은 F(A~B) = F(A) + F(B)이다.

구간 F(A~B)의 최대값은 다음과 같이 나타날 수 있다.

1. 구간 A의 최대값이 구간A~B의 최대값이다. (구간 B의 최대값이 구간 A~B의 최대값이다.)

그러면 자식노드의 구간을 그대로 할당하면 된다.

2. 구간 A~B의 최대값이 두 노드에 걸쳐있다.

부모노드C가 i~j에 대한 정보를 담고있고 (i<=j)

자식노드A,B가 각각 i~mid,mid~j에 대한 정보를 담고있다고 해보자.

우린 노드 A의 끝까지의 최대값을 알아야하고, 노드B의 시작부터의 최대값을 알아야한다.

다시 케이스를 나누자.

시작점부터의 최대값을 생각해보자.

1. 노드의 시작부터의 최대값이 왼쪽구간 자식노드의 시작부터의 최대값이다.

왼쪽 자식노드의 값을 그대로 받아오면 된다.

2. 노드의 시작부터의 최대값이 두 자식노드와 관련있다.

왼쪽구간 자식노드의 전체부분합 + 오른쪽구간 자식노드의 시작부터의 최대값을 받아오면된다.

이제 구간의 끝까지의 최대값을 생각해보자.

오른쪽 자식노드의 끝까지의 최대값을 그대로 받아오거나,

오른쪽 자석노드의 전체 구간 부분합 + 왼쪽 자식노드의 끝까지의 최대값중 더 큰것을 받아오면 된다.

즉 우리가 필요한 건 다음과 같다.

1. 구간 내 최대 부분합

2. 구간 내 시작부터의 최대값

3. 구간 내 끝까지의 최대값

4. 구간 내 전체 부분합

이를 구현하기 위해서는 설명을 따라가면 된다.

구간 내 최대부분합을 담는 세그먼트 트리를 max_tree라 하자.

구간 내 시작부터의 최대값을 담는 세그먼트 트리를 smax_tree라 하자.

구간 내 끝까지의 최대값을 담는 세그먼트 트리를 emax_tree라 하자.

구간 내 전체 부분합을 담는 세그먼트 트리를 seg_tree라 하자.

트리의 노드 node의 왼쪽구간 자식노드는 son1, 오른쪽 구간 자식노드는 son2이다.

max_tree[node] = max(max_tree[son1], max_tree[son2], emax_tree[son1] + smax_tree[son2])

smax_tree[node] = max(smax_tree[son1],seg_tree[son1] + smax_tree[son2])

emax_tree[node] = max(emax_tree[son2],seg_tree[son2] + emax_tree[son1]) 

seg_tree[node] = seg_tree[son1] + seg_tree[son2] (그냥 세그먼트 트리로 구하면 된다.)

 

구현

개인적으로 구현 난이도가 세그먼트 트리를 알고있다면 매우 쉬운 가성비좋은 자료구조라 생각한다.

세그먼트 트리 크기가 C라한다면 4*C 배열 or C*4 배열 or 각자의 방법으로 구현하면 된다.

bottom up을 기준으로 보자.

일반 세그먼트 트리는 다음과 같다.

 

def init(tree, arr, cap):
	# arr: 원본배열,tree:세그트리 cap:트리의 크기
    for i in range(cap, cap + len(arr)):
        tree[i] = arr[i - cap]
    for i in range(cap - 1, 0, -1):
        tree[i] = tree[i << 1] + tree[i << 1 | 1]


def update(tree, idx, value, cap):
	#idx:바꾸고싶은 원본 배열의 인덱스, value: 바꾸고싶은 원본 배열의 값
    i = cap + idx
    tree[i] = value
    while i > 1:
        tree[i >> 1] = tree[i] + tree[i ^ 1]
        i = i >> 1


def query(tree, start, end, cap):
	# start~end: 알고싶은 구간
    start, end = start + cap, end + cap
    result = 0
    while start < end:
        if start & 1:
            result = result + tree[start]
            start += 1
        if not end & 1:
            result = result + tree[end]
            end -= 1
        start = start >> 1
        end = end >> 1
    if start == end:
        result = result + tree[start]
    return

(추신 query의 리턴값은 null이 아닌 result여야 합니다)

이를 금광 세그로 바꾸려면 다음과 같이 바꾸면 된다.

필자의 방식을 소개하겠다.

1. C*4 형식으로 트리를 초기화한다.

tree[node][0]: 구간 내 부분합 최대

tree[node][1]: 구간 시작부터의 최대

tree[node][2]: 구간 끝까지의 최대

tree[node][3]: 구간 전체 부분합

2. 항등원은 이렇게 만들어야한다.

일반 세그트리 문제를 풀때 최대값을 구할때는 매우 작은값으로 선언하고, 부분합은 0로 선언했다.

금광세그도 똑같다.

tree[node][0~2]: 매우 작은 값(최대값을 구하므로)

tree[node][3]: 0 (합을 구하므로)

초기화

N = int(input())
arr = [0] + list(map(int, input().split())) # 원본 배열
CAP = 1 << 18 # 트리 크기
INF = 1000 * 100000 + 1 # 매우 큰 값
tree = [[-INF,-INF,-INF,0] for _ in range(CAP << 1)]  # L,R,MID,SUM


def init(left, right, node):
    if left == right:
        for i in range(4):
            tree[node][i] = arr[left]
        return tree[node]
    mid = (left + right) // 2
    A = init(left, mid, node * 2)
    B = init(mid + 1, right, node * 2 + 1)
    tree[node][0] = max(A[0], A[3] + B[0])
    tree[node][1] = max(B[1], B[3] + A[1])
    tree[node][2] = max((A[2], B[2], A[1] + B[0]))
    tree[node][3] = A[3] + B[3]
    return tree[node]

업데이트

def update(idx, value):
    i = CAP + idx
    new = tree[i][3] + value
    tree[i][0] = new
    tree[i][1] = new
    tree[i][2] = new
    tree[i][3] = new
    i = i >> 1
    while i > 0:
        A, B = tree[i << 1], tree[(i << 1) | 1]
        tree[i][0] = max(A[0], A[3] + B[0])
        tree[i][1] = max(B[1], B[3] + A[1])
        tree[i][2] = max((A[2], B[2], A[1] + B[0]))
        tree[i][3] = A[3] + B[3]
        i = i >> 1

쿼리

def query(left, right, node, a, b):
    if b < left or a > right:
        return [-INF,-INF,-INF,0]
    if a <= left and b >= right:
        return tree[node]
    result = [-INF,-INF,-INF,0]
    mid = (left + right) // 2
    A = query(left, mid, node * 2, a, b)
    B = query(mid + 1, right, node * 2 + 1, a, b)
    result[0] = max(A[0], A[3] + B[0])
    result[1] = max(B[1], B[3] + A[1])
    result[2] = max((A[2], B[2], A[1] + B[0]))
    result[3] = A[3] + B[3]
    return result

 

더 나아가서 

bottomup의 사례를 봤듯이

기존 세그먼트 트리의 형식을 안다면

쉽게 구현할 수 있다.

트리를 선언하고 각각의 필요한 연산만 하면 된다.

재귀를 사용하는 topdown도 직접 구현해보자.

또한 구간 내 최소 부분합도 항등원과 구간에 적용하는 함수만 바꿔주면 구할 수 있다.

 

추천 문제

10167번: 금광 (acmicpc.net)

10167번: 금광

(위 문제는 시간제한으로 인해 파이썬으로 풀면 안되는 문제)

13557번: 수열과 쿼리 10 (acmicpc.net)

13557번: 수열과 쿼리 10

매우 추천하는 문제이다.

어떻게 해당 자료구조를 구현했는지 이해해보자.

그 후 귀찮은 케이스처리 노가다를 하면 된다.