포인터 기반 세그먼트 트리

이 글에서는 점 갱신과 구간 합을 지원하는 세그먼트 트리를 구현합니다. 첫 번째 구현은 포인터를 사용합니다. 이 구현은 가장 빠르거나 간결하지는 않지만, 가장 범용적이고 이해하기 쉽습니다.

참고. 거의 모든 곳에서 우리는 구간 대신 $[l, r)$로 표기하는 반개구간을 사용할 것입니다. 직관에 반하기는 하지만, 이렇게 하면 코드가 약간 단순해지며, 0부터 번호를 매기는 것과 마찬가지로 일반적으로 프로그래밍에서 좋은 관행이기도 합니다.

#메모리에 저장하기

세그먼트 트리의 각 정점은 자식에 대한 참조, 자신의 경계 $[l, r)$, 그리고 문제에 필요한 추가 정보(이 경우에는 구간의 합)를 담는 구조체가 됩니다.

정점은 구축 과정에서 특정한 순서 없이 생성할 것이므로, 자식에 대한 참조는 단순히 자식 정점을 나타내는 같은 구조체를 가리키는 포인터가 됩니다.

struct Segtree {
    int lb, rb; // левые и правые границы отрезков
    int s = 0; // сумма на текущем отрезке
    Segtree *l = 0, *r = 0; // указатели на детей (0 означает, что ребенка нет)
    
    Segtree(int lb, int rb) { /* конструктор: метод, вызывающийся при создании */ }
    void add(int k, int x) { /* отреагировать на a[k] += x */ }
    int sum(int lq, int rq) { /* вывести сумму [lq, rq) */ }
};

올림피아드 문제에서 쿼리는 흔히 «덧셈에는 0 k x 문자열, 합에는 1 l r 문자열» 형식으로 주어집니다. 이는 다음과 같이 파싱할 수 있습니다:

int n, q;
cin >> n >> q;

Segtree s(0, n);

while (q--) {
    int t, x, y;
    cin >> t >> x >> y;
    if (t)
        s.add(x, y);
    else
        cout << s.sum(x, y) << endl;
}

이제 남은 것은 실제로 구축과 쿼리를 구현하는 것입니다.

#구축

세그먼트 트리는 리프에 도달할 때까지 자식을 생성하는 재귀 생성자로 구축할 수 있습니다.

Segtree(int lb, int rb) : lb(lb), rb(rb) {
    if (lb + 1 < rb) {
        // если не лист, создаем детей
        int t = (lb + rb) / 2;
        l = new Segtree(lb, t);
        r = new Segtree(t, rb);
    }
}

처음부터 배열이 0으로 채워져 있지 않다면, 링크를 연결하는 동시에 합을 계산할 수 있습니다:

Segtree(int lb, int rb) : lb(lb), rb(rb) {
    if (lb + 1 == rb)
        s = a[lb];
    else {
        int t = (lb + rb) / 2;
        l = new Segtree(lb, t);
        r = new Segtree(t, rb);
        s = l->s + r->s;
    }
}

또는 나중에 배열을 따로 순회하면서 각 수를 하나씩 더해 줄 수도 있습니다.

#변경

더하기 쿼리의 경우, 원소 $k$에 해당하는 리프에 도달할 때까지 재귀적으로 내려가면서 거치는 모든 중간 노드에 $x$를 더합니다:

void add(int k, int x) {
    s += x;
    // проверим, есть ли у нас дети:
    if (l) {
        // если k в левой половине
        if (k < l->rb)
            l->add(k, x);
        else
            r->add(k, x);
    }
}

#합

합의 경우는 더 복잡합니다. 쿼리 구간이 노드의 구간과 어떻게 겹치는지에 따라 경우를 나누어 처리해야 합니다.

int sum(int lq, int rq) {
    if (lb >= lq && rb <= rq)
        // если мы лежим полностью в отрезке запроса, вывести сумму
        return s;
    if (max(lb, lq) >= min(rb, rq))
        // если мы не пересекаемся с отрезком запроса, вывести ноль
        return 0;
    // иначе всё сложно -- запускаемся от детей и пусть они там сами решают
    return l->sum(lq, rq) + r->sum(lq, rq);
}

#최적화

이 구현은 단순하고 확장하기 쉽지만 시간과 메모리 면에서 상당히 비효율적입니다. 그럼에도 이 구현으로 프로그래밍 올림피아드 문제의 약 90%를 통과할 수 있습니다.

비교적 쉬운 최적화로는 다음과 같은 것이 있습니다.

구간의 경계 lb와 rb를 저장하지 않고, 내려가는 동안 다시 계산할 수 있습니다.
64비트 시스템에서는 new와 포인터를 사용하는 대신 정점을 배열 / 벡터에 할당하고 그에 대한 인덱스를 사용하는 것이 더 유리합니다. 인덱스는 8바이트가 아니라 4바이트를 차지하기 때문입니다.

그러나 포인터 기반 구현이 이렇게 느린 주된 이유는 필요한 노드의 데이터를 얻기 위해 포인터를 여러 번 따라가야 하기 때문입니다. 사실 노드에 고유한 번호를 매기고 그 순서대로 배열에 배치하면 포인터를 아예 없앨 수 있습니다. 이에 대한 자세한 내용은 e-maxx와 CodeForces에서 읽어 볼 수 있습니다.