[백준] 2042번 구간 합 구하기 파이썬 해설 -세그먼트 트리

알고리즘 공부 혼자 하는 게 생각보다 어렵..지만 꾸준히 해야지!

나는 원래 C언어로 알고리즘 문제를 풀었다. 임베디드 전공이다 보니 전공과목에서 프로그래밍 언어를 C/C++ 만 배웠기 때문이다. 하지만 임베디드를 넘어서 다양한 분야에 만능 프로그래머가 되고 싶어서 파이썬을 시작하게 되었고, 요즈음은 파이썬으로만 문제를 푼다.

쓸데없는 소리를 한 이유는.. 파이썬으로 알고리즘 문제를 풀면 너무 느려서 시간 초과가 너무 자주 뜬다..

편법으로는 PyPy3로 채점하는 방법이 있는데, 그래도 찝찝하긴 하다..

 

나는 이 문제를 세그먼트 트리 알고리즘으로 풀었다.

간단하게 정리해보자

 

세그먼트 트리

---

이 문제를 단순하게 배열을 확인하는 알고리즘으로 구현한다면 O(N^2)의 시간 복잡도를 갖게 되어 시간 초과를 면하지 못할 것이다.

짧게 정리하고 넘어가자면, 세그먼트 트리는 공간 복잡도는 두배로 늘어나지만, 시간 복잡도가 O(logN)로 줄어든다.

구글에 세그먼트 트리라고 검색하면 엄청 많은 자료가 나올 테니 흐름을 먼저 공부하고 코드로 구현하는 것을 추천한다.

 

 

문제 풀이

---

1. 세그먼트 트리를 초기화.

def init():
    global size, arr
    for i in range(size - 1, 0, -1):
        arr[i] = arr[i*2] + arr[i*2 + 1]

N, M, K = map(int,input().split())

size =  2**ceil(log(N,2))
size_max = size * 2
arr = [0]*(size_max)
for i in range(N):
    arr[size+i]=int(input())
init()

먼저 완성된 트리가 포화 이진트리 형태가 되기 위해 size_max를 설정해주고 배열을 초기화해 준다.

그 후 초기화된 배열의 중간 지점부터 입력을 받고, init() 함수로 앞자리까지 계산해준다.

 

2. 세그먼트 트리의 값 갱신

def modify(i, val):
    global arr
    i += size - 1
    arr[i] = val
    while i > 1:
        i //= 2
        arr[i] = arr[i*2] + arr[i*2+1]

갱신할 인덱스와 값을 입력을 받아 세그먼트 트리의 값을 갱신한다.

먼저 리프 노드들을 갱신한 후 영향을 받는 노드들을 순선대로 갱신해준다.

 

3. 세그먼트 트리에서 부분 합 구하기

def sum(left, right, node_left, node_right, node_num):
    if left > node_right or right < node_left: return 0
    if left <= node_left and right >= node_right: return arr[node_num]

    mid = (node_left + node_right)//2
    return sum(left, right, node_left, mid, node_num*2) +\
    	sum(left, right, mid+1, node_right, node_num*2 + 1)

left, right는 원하는 부분 합의 인덱스이다.

node_left, node_right는 현재 sum 함수가 순회하는 노드의 범위이다.(부분합의 범위?)

node_num은 현재 노드의 번호이다. 루트 노드를 1을 기준으로 한다.

 

결과

---

파이썬으로 정답을 제출하신 분들의 코드를 참고해보니 펜윅트리를 많이 사용하신 것 같다.

시간 복잡도는 동일하고 공간 복잡도만 줄어든다고 생각했는데, 자세한 건 알고리즘에 대해서 공부를 조금 더 해봐야 할 것 같다.

 

소스코드는 여기를 참고!

빨리 다음 문제를 풀러 가봐야겠다.