점 갱신과 구간 합 · 알고리즘 노트

개념

세그먼트 트리는 배열의 구간 정보를 트리에 저장합니다. 리프에는 원래 배열의 값을 두고, 부모에는 두 자식 구간의 합을 저장합니다. 루트에는 전체 합을 저장합니다.

배열 $[2,4,1,5]$ 의 구간 합을 저장한 트리입니다. 리프에는 배열의 값이 저장되고, 두 원소씩 묶은 구간의 합은 각각 $6$ 입니다. 루트에는 전체 합 $12$ 가 저장됩니다.

▲ [그림 1] 세그먼트 트리의 값 갱신 — 연산 전

인덱스 $2$ 의 값을 $3$ 으로 바꾸면 해당 리프의 값이 $1$ 에서 $3$ 으로 바뀝니다. 이 리프를 포함하는 오른쪽 구간의 합은 $8$ , 전체 합은 $14$ 로 갱신됩니다.

▲ [그림 2] 세그먼트 트리의 값 갱신 — 연산 후

\operatorname{tree}[i]=\operatorname{tree}[2i]+\operatorname{tree}[2i+1].

원소 수가 $N$ 이면 트리 배열은 $4N$ 칸으로 잡습니다. 루트는 t[1], 노드 t[k]의 왼쪽 자식과 오른쪽 자식은 각각 t[2 * k], t[2 * k + 1]입니다.

각 노드는 구간 $[s,e)$ 를 담당합니다. 중간 위치 $m$ 을 기준으로 왼쪽 자식은 $[s,m)$ , 오른쪽 자식은 $[m,e)$ 를 담당하며, 원소가 하나인 구간이 리프입니다.

점 갱신

점 갱신은 바꿀 인덱스가 포함된 자식으로 내려갑니다. 리프에 도착하면 값을 바꾸고, 재귀 호출에서 돌아오면서 두 자식의 합으로 부모 값을 갱신합니다.

구간 합

구간 합 질의는 $l$ 을 포함하고 $r$ 직전까지인 반열린 구간 $[l,r)$ 의 합을 구합니다. 각 노드가 담당하는 구간과 질의 구간을 비교합니다.

두 구간이 겹치지 않으면 $0$ 을 반환합니다.
노드의 구간이 질의 구간에 완전히 포함되면 저장된 합을 반환합니다.
일부만 겹치면 두 자식을 재귀 호출하고 반환값을 더합니다.

빈 구간 $[l,l)$ 의 합은 $0$ 입니다. 트리 저장 공간은 $O(N)$ 입니다.

연산	시간복잡도
점 갱신으로 원소 $N$ 개 저장	$O(N\log N)$
원소 하나 대입	$O(\log N)$
구간 합 질의	$O(\log N)$

예제

문제

배열에 다음 명령을 순서대로 수행하여라.

명령	동작
`set i x`	인덱스 $i$ 의 값을 $x$ 로 바꾼다.
`sum l r`	인덱스 $l$ 부터 $r$ 직전까지의 합을 출력한다.

입력

첫째 줄에 배열 길이 $N$ 과 명령 수 $Q$ 가 공백을 사이에 두고 주어진다. $1\le N,Q\le10\,000$ 이다.

둘째 줄에 배열의 원소 $N$ 개가 공백을 사이에 두고 주어진다. 각 원소는 $10\,000$ 보다 작은 양의 정수이다.

다음 $Q$ 개 줄에 위 표의 형식으로 명령이 하나씩 주어진다. 명령과 인자는 공백을 사이에 두고 주어진다. 인덱스와 갱신값의 범위는 $0\le i<N$ ; $0\le l\le r\le N$ ; $1\le x<10\,000$ 이다.

출력

sum 명령마다 구간의 합을 한 줄에 출력한다. $l=r$ 이면 $0$ 을 출력한다.

예제 코드

C++20

#include <iostream>
#include <string>

using namespace std;

const int N = 10000; // 최대 원소 수
int t[N * 4];

// 노드 k의 담당 구간 [s, e)에서 인덱스 i의 값을 x로 갱신
void update(int k, int s, int e, int i, int x) {
    int m;

    // 원소 하나를 담당하는 리프 노드 갱신
    if (e - s == 1) {
        t[k] = x;
        return;
    }

    m = (s + e) / 2;
    if (i < m) update(k * 2, s, m, i, x);
    else update(k * 2 + 1, m, e, i, x);

    // 자식 구간의 합으로 부모 구간의 합 갱신
    t[k] = t[k * 2] + t[k * 2 + 1];
}

// [l, r) 구간 합 질의
int query(int k, int s, int e, int l, int r) {
    int m;

    // 질의 구간과 겹치지 않으면 합에 영향을 주지 않는 0 반환
    if (r <= s || e <= l) return 0;
    // 담당 구간이 질의 구간에 포함되면 저장된 합을 바로 반환
    if (l <= s && e <= r) return t[k];

    m = (s + e) / 2;
    return query(k * 2, s, m, l, r) + query(k * 2 + 1, m, e, l, r);
}

int main() {
    int n, q, i, j, l, r, x;
    string op;

    cin >> n >> q;

    // 루트 1에서 [0, n) 구간을 기준으로 입력값 저장
    for (i = 0; i < n; ++i) {
        cin >> x;
        update(1, 0, n, i, x);
    }

    for (i = 0; i < q; ++i) {
        cin >> op;

        if (op == "set") {
            cin >> j >> x;

            update(1, 0, n, j, x);
        }
        else {
            cin >> l >> r;

            cout << query(1, 0, n, l, r) << '\n';
        }
    }
}

Java 17

import java.util.*;

public class Main {
    static final int N = 10000; // 최대 원소 수
    static int[] t = new int[N * 4];

    // 노드 k의 담당 구간 [s, e)에서 인덱스 i의 값을 x로 갱신
    static void update(int k, int s, int e, int i, int x) {
        int m;

        // 원소 하나를 담당하는 리프 노드 갱신
        if (e - s == 1) {
            t[k] = x;
            return;
        }

        m = (s + e) / 2;
        if (i < m) update(k * 2, s, m, i, x);
        else update(k * 2 + 1, m, e, i, x);

        // 자식 구간의 합으로 부모 구간의 합 갱신
        t[k] = t[k * 2] + t[k * 2 + 1];
    }

    // [l, r) 구간 합 질의
    static int query(int k, int s, int e, int l, int r) {
        int m;

        // 질의 구간과 겹치지 않으면 합에 영향을 주지 않는 0 반환
        if (r <= s || e <= l) return 0;
        // 담당 구간이 질의 구간에 포함되면 저장된 합을 바로 반환
        if (l <= s && e <= r) return t[k];

        m = (s + e) / 2;
        return query(k * 2, s, m, l, r) + query(k * 2 + 1, m, e, l, r);
    }

    public static void main(String[] args) {
        Scanner in;
        int n, q, i, j, l, r, x;
        String op;

        in = new Scanner(System.in);
        n = in.nextInt();
        q = in.nextInt();

        // 루트 1에서 [0, n) 구간을 기준으로 입력값 저장
        for (i = 0; i < n; ++i) {
            x = in.nextInt();
            update(1, 0, n, i, x);
        }

        for (i = 0; i < q; ++i) {
            op = in.next();

            if (op.equals("set")) {
                j = in.nextInt();
                x = in.nextInt();

                update(1, 0, n, j, x);
            }
            else {
                l = in.nextInt();
                r = in.nextInt();

                System.out.println(query(1, 0, n, l, r));
            }
        }
    }
}

Python 3

N = 10000  # 최대 원소 수
t = [0] * (N * 4)


# 노드 k의 담당 구간 [s, e)에서 인덱스 i의 값을 x로 갱신
def update(k, s, e, i, x):
    # 원소 하나를 담당하는 리프 노드 갱신
    if e - s == 1:
        t[k] = x
        return

    m = (s + e) // 2
    if i < m:
        update(k * 2, s, m, i, x)
    else:
        update(k * 2 + 1, m, e, i, x)

    # 자식 구간의 합으로 부모 구간의 합 갱신
    t[k] = t[k * 2] + t[k * 2 + 1]


# [l, r) 구간 합 질의
def query(k, s, e, l, r):
    # 질의 구간과 겹치지 않으면 합에 영향을 주지 않는 0 반환
    if r <= s or e <= l:
        return 0
    # 담당 구간이 질의 구간에 포함되면 저장된 합을 바로 반환
    if l <= s and e <= r:
        return t[k]

    m = (s + e) // 2
    return query(k * 2, s, m, l, r) + query(k * 2 + 1, m, e, l, r)


n, q = map(int, input().split())
lst = list(map(int, input().split()))

# 루트 1에서 [0, n) 구간을 기준으로 입력값 저장
for i, x in enumerate(lst):
    update(1, 0, n, i, x)

for _ in range(q):
    cmd = input().split()
    op = cmd[0]

    if op == "set":
        i = int(cmd[1])
        x = int(cmd[2])

        update(1, 0, n, i, x)
    else:
        l = int(cmd[1])
        r = int(cmd[2])

        print(query(1, 0, n, l, r))

1.6.4 Segment Tree

개념

점 갱신

구간 합

예제

문제

입력

출력

예제 입력

예제 출력

예제 코드