Efficient and easy segment trees에서 소개하는 일반적인 세그먼트 트리 구현보다 조금 더 빠르고, 메모리 효율성이 좋은 구현 방법을 정리하고자 한다. C++코드는 링크를 타고 확인할 수 있기 때문에 파이썬으로 정리했다. 정확히 증명을 했다기 보다는 직관적으로 이해한 느낌이라 설명이 부실할 수도 있다.


특징

  1. 재귀가 아닌 반복으로 구현한다. (파이썬에서 특히 유용하다.)
  2. 배열을 통해서 트리를 구현한다.
  3. 2^N-1 보다 적은 2N의 공간복잡도로 구현이 가능하다.
  4. 항상 완전 이진트리지만 항상 포화인 것은 아니다.
  5. 인덱스 연산에 사칙 연산 대신 비트 연산을 사용한다.
  6. query(l,r) 연산은[l, r)의 범위로 계산된다.

코드

   우선, 전체 코드를 살펴보자. 구간 합에 대한 세그먼트 트리의 생성, 수정, 질의 코드는 다음과 같다.

# 구간의 합을 구하는 세그먼트 트리
def build(arr, n):
    seg = [0] * n + arr

    for i in reversed(range(1, n)):
        seg[i] = seg[i<<1] + seg[i<<1|1]
    
    return seg
 
def modify(seg,n, idx, val):
    idx += n
    seg[idx] = val
    while idx > 1:
        seg[idx>>1] = seg[idx] + seg[idx^1]
        idx >>= 1

def query(seg, n, l ,r):
    l += n
    r += n
    result = 0
    while l<r:
        if l&1:
            result += seg[l]
            l += 1
        if r&1:
            r -= 1
            result += seg[r]
        
        l >>=1
        r >>=1
    
    return result

각 기능별로 쪼개서 하나씩 살펴보도록 하자.


build

# arr : 트리에 들어갈 원소들  ex) arr = {1,2,3,4,5}
# n   : arr의 원소의 개수    ex) n = 5
def build(arr, n):
    # 원소들이 값이 0이고 크기가 n인 배열을 만들고, 
    # 거기에 arr을 이어 붙인다.
    # ex) seg = {0,0,0,0,0,1,2,3,4,5}
    seg = [0] * n + arr
    
    # 0 이었던 n-1부터 왼쪽 자식과 오른쪽 자식의 합으로 원소를 채운다.
    # ex) seg = {0, 15, 10, 5, 9, 1, 2, 3, 4, 5}
    for i in reversed(range(1, n)):
        # 왼쪽 오른쪽 자식의 합을 저장.
        # i<<1은 항상 짝수이므로 i<<1|1 == i<<1+1
        seg[i] = seg[i<<1] + seg[i<<1|1]    
    
    return seg

   먼저, 들어오는 배열 arr의 원소의 수 만큼 트리를 위한 배열을 만들고, 배열 arr을 합쳐 총 2*N 크기의 배열을 생성한다. 이렇게 하면, 항상 잎들은 arr의 원소 값으로 채워지고, 나머지는 0인 트리를 생성할 수 있다.

   그 다음, 잎이 아닌 1~n-1번 까지의 노드(0번은 사용하지 않는다.)를 역순인 n-1부터 거꾸로 돌면서 자식들의 합으로 채워 기본적인 세그먼트 트리를 완성시킨다. 인덱스 연산에서 비트 연산으로 or을 사용하는데, 이는 인덱스에 +1을 해주는 것과 같다.


modify

# seg: 트리 배열
# n : 트리를 이루는 잎의 수 (최초 입력된 노드의 수)
# idx : 수정할 노드의 위치 (arr에서의 위치)
# val : 덮어 쓸 값.
def modify(seg,n, idx, val):
    # 1 ~ n-1은 부분합을 저장하기 때문에,
    # 실제 트리에 수정할 잎이 존재하는 위치를 n을 더해서 얻는다.
    idx += n

    # 값을 수정한다.
    seg[idx] = val
    # 구간합 노드들에 수정된 값을 반영한다.
    while idx > 1:
        # idx가 짝수면 idx^1은 idx+1
        # idx가 홀수이면 idx-1
        seg[idx>>1] = seg[idx] + seg[idx^1]
        idx >>= 1       # idx = idx // 2

   수정할 요소들은 항상 잎에 있다. 그 외의 노드들은 어떤 구간의 합을 저장하는 노드들이다. 또, 잎의 시작 위치는 기존에 n 개의 0으로 초기화 했던 것을 기억하면, n번 노드부터 잎이라는 것을 알 수 있다. 따라서, 입력된 arr의 인덱스에 n을 더해 수정해야할 요소를 찾고, 값을 val로 바꾼다.

   그 다음, root를 향해 올라가면서 변경된 값을 부분합이 저장된 중간 노드들에도 반영한다. 여기서는 xor 비트 연산을 사용했는데, 이는 idx가 부모 기준으로 왼쪽인지 오른쪽인지 알수 없기 때문이다. 비트 연산을 통해 어느쪽이든 상관 없게 간결하게 일관된 방법을 이용할 수 있다.

   이진 트리에 성질에 의해 왼쪽 자식은 짝수, 오른쪽 자식은 홀수(왼쪽 +1)로 정해져 있다. 여기서 idx가 홀수이면 이진수로 표현했을 때, 마지막 비트가 1인데, 1과 xor연산을 하면 나머지 비트는 그대로고, 마지막 비트가 0으로 변해 -1 한 것과 같은 효과가 있다. 반대로, idx가 짝수이면 마지막 비트가 0이고, 1과 xor 연산을 하면 나머지 비트는 그대로고, 마지막 비트가 1로 바뀌어 +1과 같은 효과가 있다.


query

# seg : 트리가 저장된 배열
# n : 트리의 잎 수
# l : 왼쪽 범위
# r : 오른쪽 범위
def query(seg, n, l ,r):
    l += n          # 왼쪽 구간의 잎의 위치
    r += n          # 오른쪽 구간의 잎의 위치
    result = 0      # 질의의 결과
    while l<r:
        # L이 홀수(right)이면 자신을 더하고, 옆 칸으로 넘어간다.
        if l&1:
            result += seg[l]
            l += 1

        # R이 홀수(right)이면 자신의 형제 (같은 부모의 왼쪽 자식)을 더한다.
        if r&1:     
            r -= 1
            result += seg[r]
        
        # 각각 부모로 이동
        l >>=1
        r >>=1
    
    return result

   전혀 직관적이지 않은 코드다. 특이한 점은 매 반복마다 홀수 일 때만 어떤 액션을 취해주는데, 왼쪽에서는 인덱스가 홀수인 자기 자신을 더하고, 오른쪽 구간에서는 항상 옆에 있는 같은 부모의 왼쪽 형제를 더한다는 것이다. 각각의 구간이 홀수, 짝수일 때를 각각 따져보자.

   우선, 왼쪽 닫힌 구간이 홀수이면 어떤 노드의 오른쪽 자식이라는 말이 된다. 그 말은 왼쪽 형제는 범위에 포함되지 않는다는 소리이고, 왼쪽 자식의 값이 포함된 부모 역시 범위에 포함될 수 없다. 따라서, 자기 자신을 더하고, 1을 더해 범위를 한칸 이동시킨다.(범위를 좁힌다.) 여기서 1을 더해도 문제가 없는 이유는 같은 레벨(높이)에서 오른쪽은 항상 범위에 포함되기 때문이다.

   반대로 왼쪽 닫힌 구간이 짝수라면, 어떤 중간 노드의 왼쪽 자식이 된다. 이 경우에는 양쪽 형제들이 모두 범위에 포함되고, 그 값은 부모가 가지고 있기 때문에 부모쪽에서 계산하도록 아무 행동도 하지 않는다.

   오른쪽 열린 구간이 홀수이면 역시 어떤 노드의 오른쪽 자식이고, 항상 왼쪽 형제가 존재한다. 열린 구간이기 때문에 자신은 포함되어서는 안되지만, 형제는 범위에 포함되기 때문에 -1을 해 형제로 이동하고, 형제를 결과에 더한다.

   마지막으로 오른쪽 열린 구간이 짝수이면, 이쪽 트리는 구간에 포함되지 않는다고 볼 수 있다. 따라서, 아무 행동도 하지 않는다.


예시

   글로만 이루어진 설명으로는 이해하기 힘들어 예제를 준비했다. 기존에 예를 들었던 배열 arr=[1,2,3,4,5]로 트리를 생성하고, qeury(seg,5, 0, 4)라는 연산을 수행한다고 하자. arr의 인덱스 0, 1, 2, 3의 부분합인 1+2+3+4=10이 나와야 한다. 아래 트리 사진으로 더할 구간을 눈으로 확인할 수 있다.


   코드를 따라가서, 가장 먼저 두 값에 n을 더해 잎의 위치를 찾으면 왼쪽 닫힌 구간(L)과 오른쪽 열린 구간(R)은 각각 5와 9가 된다. 두 구간 다 홀수이기 때문에 if에 걸린다. 그러면 seg[5]와 seg[8]의 값이 더해지고 범위는 6과 8이 된다. 반복의 마지막 연산으로 부모로 이동하면 3과 4가 된다.


   그 다음 반복에서, 3과 4 중 왼쪽 닫힌 구간(L)은 자기의 범위를 포함하기 때문에 자신을 더한 다음, 4로 이동하고 부모인 2로 이동한다. 오른쪽 열린 구간은 짝수기 때문에 바로 부모인 2로 올라간다.


   두 구간의 값이 같아 졌기 때문에 더 이상의 반복을 진행하지 않는다. 3번 5번 8번 인덱스를 더함으로 원하던 [0, 4) 구간의 합 10을 구할 수 있다.


기타

   해당 자료구조의 구현이 삽입 연산, Lazy Propagation 등 역시 구현할 수 있다고 한다. 그러나, 이러한 문제를 아직 풀지 않아서, 나중에 해당 문제들을 이해하게 된다면 추가로 정리하도록 하겠다.