[Python] 구간 최댓값 문제 풀이

구간 최댓값과 세그먼트 트리

잊을만하면 나오는 "구간 최댓값" 문제. 

구간 최댓값, 구간합, 구간 최솟값 같은 문제는 세그먼트 트리 방식을 사용하는 것이 시간 복잡도 측면에서 유리할 때가 많다.

 

 

이를 트리 구조로 다시 봐보자.

 

세그먼트 트리이지만, 개인적으로는 트리가 아니라 리스트 방식으로 보는 것이 이해가 훨씬 더 수월했다. 보는 바와 같이 인덱스 1부터가 사실상 시작인데, 이는 트리 구조에서 인덱스 1을 root node로 놓으면, 2x하면 왼쪽 자식, 2x+1 하면 오른쪽 자식이 되어 계산이 편해지기 때문이다.

 

그럼 [1, 3, 5, 7, 9, 11] 중 index range(1, 4), 즉, [3, 5, 7] 의 최댓값을 찾아보자.

1. 우선 우리가 구하고자 하는 구간 최대값 \(x = - \inf \) 라고 하자.

2. index의 \(l = 1, r = 4\) 이다.

3. 이것들에 6을 더하자. \(l = 7, r = 10\) 이 된다.

4. \(l<r\) 인 동안, 

    i. \(l \) 이 홀수라면, \( x = \max (x, T_l), l \leftarrow l+1 \)

    ii. \(r\) 이 홀수라면, \( r \leftarrow r-1 ,x = \max (x, T_r)\)

    iii. \( l \leftarrow \lfloor l / 2 \rfloor \)

    iv. \( r \leftarrow \lfloor r / 2 \rfloor \)

 

 

그러므로 다음과 같이 된다.

 

1단계. \(l = 7, r = 10, x = - \inf \)

2단계. \(l = 8, r = 10, x = 3\)

3단계. \(l = 4, r = 4, x = 7\)

끝! 

 

구간 최댓값 7을 잘 찾아내는 것을 볼 수 있다.

 

코드를 보자. 생각보다 간단하다.

class SegmentTree:
    def __init__(self, data):
        self.n = len(data)
        self.tree = [0] * (2 * self.n)

        for i in range(self.n):
            self.tree[self.n + i] = data[i]

        for i in range(self.n - 1, 0, -1):
            self.tree[i] = max(self.tree[2 * i], self.tree[2 * i + 1])

    def find_max(self, l, r):
        max_value = float('-inf')
        l += self.n
        r += self.n

        while l < r:
            if l % 2 == 1:
                max_value = max(max_value, self.tree[l])
                l += 1

            if r % 2 == 1:
                r -= 1
                max_value = max(max_value, self.tree[r])

            l //= 2
            r //= 2

        return max_value