구간 쿼리와 누적합을 위한 트리: 세그먼트 트리, 펜윅 트리, 트라이 구현
들어가며: “구간 합 쿼리가 O(n)이라 느려요”
길이 n인 배열에서 구간 [L, R]의 합을 q번 묻는데, 쿼리마다 구간을 처음부터 순회하면 전체 작업량은 O(nq)입니다. n과 q가 모두 10만이면 최악의 경우 100억 번의 덧셈이 필요하고, 코딩 테스트의 일반적인 시간 제한을 훌쩍 넘습니다. 업데이트가 없다면 누적합 배열 하나로 쿼리를 O(1)에 끝낼 수 있지만, 원소가 계속 바뀌면 누적합을 매번 다시 계산해야 해서 이번에는 업데이트가 O(n)이 됩니다.
세그먼트 트리와 펜윅 트리는 이 둘 사이의 균형점을 찾습니다. 쿼리와 업데이트를 모두 O(log n)에 처리하므로, 같은 조건에서 연산 수가 약 10만 × 17 ≈ 170만 번 수준으로 줄어듭니다.
flowchart TD
subgraph wrong["단순 반복: 쿼리당 O(n)"]
W1[쿼리 10만 번] --> W2["매번 [L,R] 순회"]
W2 --> W3["최악 10만 × 10만 = 100억"]
end
subgraph right["세그먼트/펜윅 트리: 쿼리당 O(log n)"]
R1[쿼리 10만 번] --> R2["트리에서 구간 합"]
R2 --> R3["10만 × log₂(10만) ≈ 170만"]
end
이 글에서는 구간 쿼리와 접두사 검색이 필요한 상황을 먼저 정리하고, 세그먼트 트리·펜윅 트리·트라이를 C++로 구현한 뒤, 구현하면서 자주 틀리는 부분을 짚습니다. 안전한 쿼리 예제가 std::optional을 쓰므로 C++17 이상으로 컴파일합니다.
어떤 문제에 어떤 트리를 쓰는가
구간 최솟값을 자주 묻고 값도 자주 바뀌는 경우를 생각해 봅시다. 예를 들어 시계열 가격 데이터에서 “3일부터 7일까지의 최저가”를 반복해서 조회하면서, 새 가격이 들어올 때마다 값을 갱신해야 한다면 쿼리마다 순회하는 방식은 곧바로 한계에 부딪힙니다.
// 쿼리마다 O(n) 순회
std::vector<int> prices = {100, 95, 102, 98, 105, 103, 107};
int query_min(int L, int R) {
int mn = INT_MAX;
for (int i = L; i <= R; ++i) mn = std::min(mn, prices[i]);
return mn;
}최솟값은 누적합처럼 “앞에서부터의 결과 두 개를 빼서” 구간 값을 얻을 수 없으므로, 세그먼트 트리가 필요한 대표적인 경우입니다.
반대로 합만 필요하다면 누적합 배열을 떠올리게 됩니다. 문제는 원소 하나가 바뀌면 그 뒤의 누적합이 모두 바뀐다는 점입니다.
// 누적합 배열은 업데이트 한 번에 O(n)
void update(int i, int delta) {
arr[i] += delta;
for (int j = i; j < n; ++j) prefix[j] += delta;
}이 경우에는 펜윅 트리(Binary Indexed Tree)가 코드도 짧고 메모리도 적게 씁니다.
자동완성처럼 접두사로 문자열을 찾는 문제는 성격이 다릅니다. 사실 정렬된 std::map이나 정렬된 std::vector에서 lower_bound("app")로 시작 위치를 찾고, 접두사가 맞지 않을 때까지 순회하는 방식도 O(m log n + k)로 충분히 빠릅니다(m은 접두사 길이, k는 결과 개수).
std::map<std::string, int> words;
for (auto it = words.lower_bound("app"); it != words.end(); ++it) {
if (it->first.compare(0, 3, "app") != 0) break;
suggest(it->first);
}트라이가 유리한 경우는 공통 접두사가 많은 대량의 단어를 저장해 접두사를 공유하고 싶을 때, 한 글자씩 입력될 때마다 이전 노드에서 이어서 탐색하고 싶을 때, 그리고 노드마다 부가 정보(예: 서브트리의 인기 단어 상위 목록)를 붙이고 싶을 때입니다.
순위 쿼리도 세그먼트 트리나 펜윅 트리로 풀 수 있습니다. 점수 값의 범위를 인덱스로 삼아 “그 점수를 가진 사용자 수”를 저장해 두면, “누적 개수가 k 이상이 되는 가장 작은 인덱스”가 k번째로 작은 점수입니다. 누적합에 이분 탐색을 씌우면 O(log² n)이고, 펜윅 트리의 이진 리프팅이나 세그먼트 트리에서 루트부터 내려가는 방식을 쓰면 O(log n)에 찾을 수 있습니다. 매번 정렬하는 O(n log n) 방식과 비교하면 차이가 큽니다.
| 문제 유형 | 추천 자료구조 | 쿼리 | 업데이트 | 비고 |
|---|---|---|---|---|
| 구간 합, 구간 XOR | 펜윅 트리 | O(log n) | O(log n) | 역연산(뺄셈, XOR)이 있음 |
| 구간 최소/최대, GCD | 세그먼트 트리 | O(log n) | O(log n) | 역연산이 없어 펜윅 부적합 |
| 구간에 값 더하기 + 구간 합 | lazy 세그먼트 트리 | O(log n) | O(log n) | 구간 업데이트 |
| 접두사 검색 | 트라이 | O(m) | O(m) | m = 문자열 길이 |
| k번째 원소 | 세그먼트 트리/펜윅 트리 | O(log n) | O(log n) | 트리를 내려가며 탐색 |
세그먼트 트리 구현
핵심 아이디어
세그먼트 트리는 각 노드가 배열의 한 구간을 담당하는 이진 트리입니다. 루트는 [0, n-1], 왼쪽 자식은 [0, mid], 오른쪽 자식은 [mid+1, n-1]을 맡고, 각 노드에는 자기 구간의 합(또는 최솟값 등)을 저장합니다.
flowchart TD R["[0,6]"] --> L["[0,3]"] R --> RR["[4,6]"] L --> LL["[0,1]"] L --> LR["[2,3]"] RR --> RL["[4,5]"] RR --> RRR["[6,6]"]
쿼리 구간 [ql, qr]이 주어지면, 노드 구간이 쿼리 구간에 완전히 포함될 때는 저장된 값을 바로 쓰고, 전혀 겹치지 않으면 항등원을 돌려주며, 일부만 겹치면 두 자식으로 내려갑니다. 트리의 각 깊이에서 “일부만 겹치는” 노드는 최대 두 개뿐이므로 방문 노드 수가 O(log n)으로 제한됩니다. 노드 번호는 루트를 1로 두고 노드 i의 자식을 2i, 2i+1로 매기며, 배열은 4n 크기면 충분합니다.
구간 합 세그먼트 트리
#include <cassert>
#include <vector>
// 0-indexed 구간 합 세그먼트 트리
// build O(n), query O(log n), update O(log n)
class SegmentTree {
std::vector<long long> tree_;
int n_;
long long queryImpl(int node, int l, int r, int ql, int qr) const {
if (qr < l || r < ql) return 0; // 겹치지 않음
if (ql <= l && r <= qr) return tree_[node]; // 완전히 포함
int mid = (l + r) / 2;
return queryImpl(2 * node, l, mid, ql, qr)
+ queryImpl(2 * node + 1, mid + 1, r, ql, qr);
}
void updateImpl(int node, int l, int r, int idx, long long delta) {
if (l == r) {
tree_[node] += delta;
return;
}
int mid = (l + r) / 2;
if (idx <= mid) updateImpl(2 * node, l, mid, idx, delta);
else updateImpl(2 * node + 1, mid + 1, r, idx, delta);
tree_[node] = tree_[2 * node] + tree_[2 * node + 1];
}
void build(int node, int l, int r, const std::vector<long long>& arr) {
if (l == r) {
tree_[node] = arr[l];
return;
}
int mid = (l + r) / 2;
build(2 * node, l, mid, arr);
build(2 * node + 1, mid + 1, r, arr);
tree_[node] = tree_[2 * node] + tree_[2 * node + 1];
}
public:
explicit SegmentTree(const std::vector<long long>& arr)
: tree_(4 * arr.size()), n_(static_cast<int>(arr.size())) {
if (n_ > 0) build(1, 0, n_ - 1, arr);
}
long long query(int ql, int qr) const {
assert(0 <= ql && ql <= qr && qr < n_);
return queryImpl(1, 0, n_ - 1, ql, qr);
}
// idx번 원소에 delta를 더한다
void update(int idx, long long delta) {
assert(0 <= idx && idx < n_);
updateImpl(1, 0, n_ - 1, idx, delta);
}
};업데이트는 대상 인덱스가 속한 쪽 자식으로만 내려가면 되므로, 양쪽을 모두 호출하고 범위 밖에서 돌아오게 하는 것보다 idx <= mid로 분기하는 편이 호출 수가 절반입니다. 원소 수가 0일 때 build(1, 0, -1)이 호출되지 않도록 생성자에서 막아 두었습니다.
구간 최솟값 세그먼트 트리 (RMQ)
합 대신 최솟값을 저장하고, 겹치지 않는 구간에서는 항등원인 INT_MAX를 돌려주면 됩니다. 업데이트는 델타를 더하는 대신 값을 교체합니다.
#include <algorithm>
#include <climits>
#include <vector>
class SegmentTreeMin {
std::vector<int> tree_;
int n_;
static constexpr int INF = INT_MAX;
int queryImpl(int node, int l, int r, int ql, int qr) const {
if (qr < l || r < ql) return INF;
if (ql <= l && r <= qr) return tree_[node];
int mid = (l + r) / 2;
return std::min(queryImpl(2 * node, l, mid, ql, qr),
queryImpl(2 * node + 1, mid + 1, r, ql, qr));
}
void updateImpl(int node, int l, int r, int idx, int val) {
if (l == r) {
tree_[node] = val;
return;
}
int mid = (l + r) / 2;
if (idx <= mid) updateImpl(2 * node, l, mid, idx, val);
else updateImpl(2 * node + 1, mid + 1, r, idx, val);
tree_[node] = std::min(tree_[2 * node], tree_[2 * node + 1]);
}
void build(int node, int l, int r, const std::vector<int>& arr) {
if (l == r) {
tree_[node] = arr[l];
return;
}
int mid = (l + r) / 2;
build(2 * node, l, mid, arr);
build(2 * node + 1, mid + 1, r, arr);
tree_[node] = std::min(tree_[2 * node], tree_[2 * node + 1]);
}
public:
explicit SegmentTreeMin(const std::vector<int>& arr)
: tree_(4 * arr.size(), INF), n_(static_cast<int>(arr.size())) {
if (n_ > 0) build(1, 0, n_ - 1, arr);
}
int query(int ql, int qr) const { return queryImpl(1, 0, n_ - 1, ql, qr); }
void update(int idx, int val) { updateImpl(1, 0, n_ - 1, idx, val); }
};#include <iostream>
int main() {
std::vector<long long> arr = {1, 3, 5, 7, 9, 11};
SegmentTree st(arr);
std::cout << st.query(1, 3) << "\n"; // 3+5+7 = 15
st.update(2, 10); // arr[2]: 5 → 15
std::cout << st.query(1, 3) << "\n"; // 3+15+7 = 25
std::vector<int> mins = {2, 5, 1, 8, 3, 9};
SegmentTreeMin stMin(mins);
std::cout << stMin.query(0, 2) << "\n"; // min(2,5,1) = 1
}예제: 값 교체 업데이트와 구간 합 (LeetCode 307)
LeetCode 307 “Range Sum Query - Mutable”은 update(index, val)이 델타가 아니라 새 값을 받습니다. 원본 배열을 함께 보관해 두고 val - nums[index]를 델타로 바꿔 넘기면 앞의 SegmentTree를 그대로 쓸 수 있습니다.
class NumArray {
std::vector<int> nums_;
SegmentTree st_;
static std::vector<long long> widen(const std::vector<int>& v) {
return std::vector<long long>(v.begin(), v.end());
}
public:
explicit NumArray(const std::vector<int>& nums)
: nums_(nums), st_(widen(nums)) {}
void update(int index, int val) {
st_.update(index, static_cast<long long>(val) - nums_[index]);
nums_[index] = val;
}
long long sumRange(int left, int right) const {
return st_.query(left, right);
}
};
// NumArray na({1, 3, 5}); na.sumRange(0, 2) == 9; na.update(1, 2); na.sumRange(0, 2) == 8;델타를 계산할 때 val - nums_[index]를 int끼리 빼면, 두 값이 부호가 다르고 크기가 클 때 뺄셈 자체가 오버플로할 수 있으므로 한쪽을 먼저 long long으로 바꿨습니다. 같은 문제를 아래의 펜윅 트리로도 풀 수 있는데, 인덱스를 1 올려서 넘기는 것만 다릅니다.
class NumArrayFenwick {
FenwickTree ft_;
std::vector<int> nums_;
public:
explicit NumArrayFenwick(const std::vector<int>& nums)
: ft_(static_cast<int>(nums.size())), nums_(nums) {
for (int i = 0; i < static_cast<int>(nums.size()); ++i) ft_.add(i + 1, nums[i]);
}
void update(int index, int val) {
ft_.add(index + 1, static_cast<long long>(val) - nums_[index]);
nums_[index] = val;
}
long long sumRange(int left, int right) const {
return ft_.rangeSum(left + 1, right + 1);
}
};펜윅 트리 구현
핵심 아이디어
펜윅 트리는 1-indexed 배열 tree에서 tree[i]가 [i - (i & -i) + 1, i] 구간의 합을 담도록 구성합니다. i & -i는 i의 이진 표현에서 가장 낮은 1비트의 값(lowbit)입니다. 예를 들어 6은 이진수 110이므로 lowbit가 2이고, tree[6]은 [5, 6]의 합을 담습니다.
| i | 이진수 | lowbit | 담당 구간 |
|---|---|---|---|
| 1 | 001 | 1 | [1, 1] |
| 2 | 010 | 2 | [1, 2] |
| 3 | 011 | 1 | [3, 3] |
| 4 | 100 | 4 | [1, 4] |
| 5 | 101 | 1 | [5, 5] |
| 6 | 110 | 2 | [5, 6] |
누적합 [1, i]를 구할 때는 i에서 lowbit를 계속 빼면서 해당 칸을 더하고, 원소 i를 갱신할 때는 lowbit를 계속 더하면서 i를 포함하는 칸들을 갱신합니다. 어느 쪽이든 매 단계에서 1비트가 하나씩 사라지거나 자리올림이 일어나므로 반복 횟수는 O(log n)입니다.
#include <cassert>
#include <vector>
// 1-indexed 펜윅 트리
class FenwickTree {
int n_;
std::vector<long long> tree_;
public:
explicit FenwickTree(int n) : n_(n), tree_(n + 1, 0) {}
// i번째(1-indexed) 원소에 delta를 더한다
void add(int i, long long delta) {
assert(1 <= i && i <= n_);
for (; i <= n_; i += (i & -i)) tree_[i] += delta;
}
// [1, i] 구간 합
long long prefixSum(int i) const {
assert(0 <= i && i <= n_);
long long sum = 0;
for (; i > 0; i -= (i & -i)) sum += tree_[i];
return sum;
}
// [L, R] 구간 합 (1-indexed)
long long rangeSum(int L, int R) const {
assert(1 <= L && L <= R && R <= n_);
return prefixSum(R) - prefixSum(L - 1);
}
// 0-indexed 배열로 초기화 (O(n log n))
void init(const std::vector<long long>& arr) {
for (int i = 0; i < static_cast<int>(arr.size()) && i < n_; ++i) add(i + 1, arr[i]);
}
};멤버 초기화 순서는 선언 순서를 따르므로, 생성자 초기화 목록도 선언 순서(n_, tree_)와 맞춰 두어야 -Wreorder 경고가 나지 않습니다. init은 원소마다 add를 호출해 O(n log n)인데, 각 칸의 값을 부모 칸(i + lowbit(i))에 한 번씩 더해 주는 방식으로 O(n)에 구성할 수도 있습니다.
int main() {
FenwickTree ft(6);
for (int i = 1; i <= 6; ++i) ft.add(i, 2 * i - 1); // 1 3 5 7 9 11
std::cout << ft.rangeSum(2, 4) << "\n"; // 3+5+7 = 15
ft.add(3, 10);
std::cout << ft.rangeSum(2, 4) << "\n"; // 3+15+7 = 25
}펜윅 트리와 세그먼트 트리 비교
| 항목 | 펜윅 트리 | 세그먼트 트리 |
|---|---|---|
| 구현 분량 | 핵심 루프 두 개 | 재귀 함수 세 개 |
| 메모리 | n+1칸 | 재귀 구현 기준 약 4n칸 |
| 구간 합 | O(log n) | O(log n) |
| 구간 최소/최대 | 불가 | 가능 |
| 구간 업데이트 | 차분 배열을 쓰는 변형으로 가능 | lazy propagation |
펜윅 트리가 구간 값을 구하는 방식은 prefix(R) - prefix(L-1), 즉 역연산입니다. 합은 뺄셈, XOR은 다시 XOR로 되돌릴 수 있지만, 최솟값은 min(a, b)의 결과만으로 a를 제거한 값을 복원할 수 없습니다. 그래서 구간 최솟값·최댓값·GCD는 세그먼트 트리를 써야 합니다.
트라이 구현
트라이는 문자열을 한 글자씩 따라 내려가는 트리입니다. “apple”을 넣으면 루트에서 a→p→p→l→e 경로가 생기고, 마지막 노드에 “여기서 끝나는 단어가 있다”는 표시를 합니다. 같은 접두사를 가진 단어는 경로를 공유합니다.
flowchart TD R["(루트)"] --> A[a] A --> P1[p] P1 --> P2[p] P2 --> L[l] L --> E["e (끝)"] A --> N[n] N --> T["t (끝)"]
#include <array>
#include <memory>
#include <string>
#include <vector>
// 영소문자 전용 트라이
class Trie {
struct Node {
std::array<std::unique_ptr<Node>, 26> children{};
bool is_end = false;
};
std::unique_ptr<Node> root_ = std::make_unique<Node>();
static int index(char c) {
return (c >= 'a' && c <= 'z') ? c - 'a' : -1;
}
const Node* find(const std::string& s) const {
const Node* cur = root_.get();
for (char c : s) {
int i = index(c);
if (i < 0 || !cur->children[i]) return nullptr;
cur = cur->children[i].get();
}
return cur;
}
void collect(const Node* node, std::string& path, std::vector<std::string>& out) const {
if (node->is_end) out.push_back(path);
for (int i = 0; i < 26; ++i) {
if (node->children[i]) {
path.push_back(static_cast<char>('a' + i));
collect(node->children[i].get(), path, out);
path.pop_back();
}
}
}
public:
// 허용되지 않는 문자가 있으면 아무것도 삽입하지 않고 false
bool insert(const std::string& word) {
for (char c : word) if (index(c) < 0) return false;
Node* cur = root_.get();
for (char c : word) {
int i = index(c);
if (!cur->children[i]) cur->children[i] = std::make_unique<Node>();
cur = cur->children[i].get();
}
cur->is_end = true;
return true;
}
bool search(const std::string& word) const {
const Node* n = find(word);
return n && n->is_end;
}
bool startsWith(const std::string& prefix) const { return find(prefix) != nullptr; }
// prefix로 시작하는 모든 단어 (사전순)
std::vector<std::string> wordsWithPrefix(const std::string& prefix) const {
std::vector<std::string> result;
if (const Node* n = find(prefix)) {
std::string path = prefix;
collect(n, path, result);
}
return result;
}
};원래 흔히 보는 구현은 삽입 도중 허용되지 않는 문자를 만나면 그냥 반환하는데, 그러면 그 앞까지 만든 노드가 단어 끝 표시 없이 남습니다. 위 코드는 삽입 전에 먼저 문자를 검사해서 이런 부분 경로가 생기지 않게 했습니다. collect는 prefix + c로 매번 새 문자열을 만드는 대신 하나의 path 버퍼에 글자를 넣고 빼면서 순회하므로 불필요한 할당이 줄어듭니다. 자식 배열을 0부터 25까지 순서대로 방문하므로 결과는 사전순입니다.
#include <iostream>
int main() {
Trie trie;
trie.insert("apple");
trie.insert("application");
trie.insert("apply");
trie.insert("banana");
std::cout << trie.search("apple") << "\n"; // 1
std::cout << trie.search("app") << "\n"; // 0
std::cout << trie.startsWith("app") << "\n"; // 1
for (const auto& w : trie.wordsWithPrefix("app")) std::cout << w << " ";
// apple application apply
}search와 startsWith는 O(m)이지만, wordsWithPrefix는 접두사 아래 서브트리 전체를 도므로 O(m + 서브트리 노드 수)입니다. 짧은 접두사(“a” 한 글자 등)로 부르면 사실상 사전 전체를 순회할 수 있으니, 실제 자동완성에서는 결과 개수 제한을 두거나 노드마다 상위 후보를 미리 저장해 둡니다.
노드마다 26칸짜리 포인터 배열을 두면 64비트 환경에서 노드 하나가 200바이트를 넘습니다. 대부분의 노드는 자식이 한두 개뿐이므로, 문자 집합이 크거나(유니코드) 단어가 많을 때는 자식을 정렬된 작은 벡터나 해시맵으로 두거나, 자식이 하나뿐인 경로를 한 간선으로 합치는 압축 트라이(radix tree, Patricia trie)를 고려합니다. 압축 트라이는 “application”과 “apply”가 “appl”까지 한 간선을 공유하고 거기서 갈라지는 식으로 노드 수를 줄이지만, 삽입 시 간선을 쪼개는 처리가 필요해 구현이 꽤 복잡해집니다.
구현하면서 자주 틀리는 부분
인덱스 기준 혼동
세그먼트 트리는 보통 0-indexed로, 펜윅 트리는 1-indexed로 구현합니다. 문제 입력이 1-indexed라면 세그먼트 트리에는 st.query(L - 1, R - 1)처럼 1을 빼서 넘겨야 합니다. 펜윅 트리에 add(0, x)를 호출하면 0 & -0이 0이라 인덱스가 영원히 0에 머물러 무한 루프에 빠집니다. 0-indexed 배열을 넣을 때는 ft.add(i + 1, arr[i])처럼 변환합니다. 한 코드베이스에 두 자료구조가 섞여 있다면 클래스 선언부 주석에 인덱스 기준을 적어 두는 것이 가장 확실한 예방책입니다.
트리 배열 크기
재귀로 구간을 반씩 나누는 세그먼트 트리는 n이 2의 거듭제곱이 아닐 때 트리가 한쪽으로 한 단계 더 깊어지고, 노드 번호가 2n을 넘을 수 있습니다. 필요한 최대 크기는 n보다 크거나 같은 가장 작은 2의 거듭제곱의 두 배로, 항상 4n 이하입니다. 그래서 관례적으로 4n을 할당합니다. n만 할당하면 범위 밖 쓰기가 일어나며, std::vector의 operator[]는 경계를 검사하지 않으므로 바로 크래시가 나지 않고 다른 데이터를 망가뜨릴 수 있습니다.
빈 구간과 경계
query(L, R)에 L > R이 들어오면 구현에 따라 이상한 값이 나옵니다. 빈 구간을 정상 입력으로 허용해야 한다면 호출 전에 항등원(합이면 0)을 반환하도록 처리하고, 허용하지 않는다면 위 구현처럼 assert로 막습니다. 외부 입력을 그대로 받는 경우에는 std::optional로 감싼 래퍼를 두는 편이 안전합니다.
#include <optional>
std::optional<long long> safeQuery(const SegmentTree& st, int ql, int qr, int n) {
if (ql < 0 || qr >= n || ql > qr) return std::nullopt;
return st.query(ql, qr);
}합의 오버플로
원소가 최대 10만이고 개수가 10만이면 합은 100억까지 커질 수 있어 int(약 ±21억)를 넘습니다. 흔한 실수는 반환형만 long long으로 바꾸고 tree_는 std::vector<int>로 두는 것인데, 이러면 내부 노드에 합을 저장하는 순간 이미 넘칩니다. 트리 배열의 원소 타입 자체를 long long으로 둬야 합니다.
항등원
겹치지 않는 구간에서 돌려주는 값은 병합 연산의 항등원이어야 합니다. 합과 XOR은 0, 곱은 1, 최솟값은 INT_MAX, 최댓값은 INT_MIN, GCD는 0입니다. 최솟값 트리에서 습관적으로 0을 돌려주면, 음수가 없는 배열에서 모든 쿼리 결과가 0이 됩니다.
구간 업데이트: lazy propagation
“구간 [L, R]의 모든 원소에 delta를 더하라”는 업데이트를 점 업데이트로 처리하면 O(n log n)입니다. lazy propagation은 구간 전체가 업데이트 범위에 들어가는 노드에서 멈추고, 자식에게 내려보낼 값을 lazy 배열에 적어 둡니다. 그 자식은 나중에 실제로 방문될 때 값을 반영합니다.
여기서는 lazy[node]를 “아직 이 노드의 tree_ 값에 반영되지 않은, 구간의 모든 원소에 더할 값”으로 정의합니다. 이 정의에서는 노드를 방문할 때마다 먼저 자기 lazy를 tree_[node]에 반영하고, 그 값을 자식들의 lazy로 넘긴 뒤 0으로 초기화해야 합니다. 순서가 뒤바뀌어 자식에게 넘기기 전에 0으로 지우거나, 반영하지 않고 넘기기만 하면 값이 사라지거나 두 번 더해집니다.
class LazySegmentTree {
std::vector<long long> tree_, lazy_;
int n_;
void push(int node, int l, int r) {
if (lazy_[node] == 0) return;
tree_[node] += lazy_[node] * (r - l + 1);
if (l != r) {
lazy_[2 * node] += lazy_[node];
lazy_[2 * node + 1] += lazy_[node];
}
lazy_[node] = 0;
}
void updateImpl(int node, int l, int r, int ql, int qr, long long delta) {
push(node, l, r);
if (qr < l || r < ql) return;
if (ql <= l && r <= qr) {
lazy_[node] += delta;
push(node, l, r);
return;
}
int mid = (l + r) / 2;
updateImpl(2 * node, l, mid, ql, qr, delta);
updateImpl(2 * node + 1, mid + 1, r, ql, qr, delta);
tree_[node] = tree_[2 * node] + tree_[2 * node + 1];
}
long long queryImpl(int node, int l, int r, int ql, int qr) {
push(node, l, r);
if (qr < l || r < ql) return 0;
if (ql <= l && r <= qr) return tree_[node];
int mid = (l + r) / 2;
return queryImpl(2 * node, l, mid, ql, qr)
+ queryImpl(2 * node + 1, mid + 1, r, ql, qr);
}
public:
explicit LazySegmentTree(int n) : tree_(4 * n), lazy_(4 * n), n_(n) {}
void rangeAdd(int ql, int qr, long long delta) { updateImpl(1, 0, n_ - 1, ql, qr, delta); }
long long query(int ql, int qr) { return queryImpl(1, 0, n_ - 1, ql, qr); }
};updateImpl이 일부만 겹치는 노드에서 두 자식을 모두 호출하는 데는 이유가 있습니다. 범위 밖 자식도 함수 첫 줄에서 push를 해서 자기 tree_ 값을 최신으로 만든 뒤 돌아오므로, 부모가 tree_[2*node] + tree_[2*node+1]로 다시 합칠 때 두 값이 모두 정확합니다. 범위 밖 자식을 호출하지 않도록 최적화하려면 부모에서 합칠 때 자식의 lazy를 따로 고려해야 합니다.
반복문 기반 세그먼트 트리
재귀 세그먼트 트리의 재귀 깊이는 log₂ n 정도(n이 100만이어도 약 20)라서 스택 오버플로는 걱정할 필요가 없습니다. 반복문 구현을 쓰는 이유는 함수 호출 비용이 없어 상수 배가 작고, 메모리를 2n만 쓰기 때문입니다. 리프를 tree_[n..2n-1]에 두고 노드 i의 부모를 i/2로 두는 레이아웃입니다.
class SegmentTreeIter {
int n_;
std::vector<long long> tree_;
public:
explicit SegmentTreeIter(const std::vector<long long>& arr)
: n_(static_cast<int>(arr.size())), tree_(2 * arr.size()) {
for (int i = 0; i < n_; ++i) tree_[n_ + i] = arr[i];
for (int i = n_ - 1; i > 0; --i) tree_[i] = tree_[2 * i] + tree_[2 * i + 1];
}
void update(int idx, long long delta) {
for (idx += n_; idx > 0; idx /= 2) tree_[idx] += delta;
}
// [ql, qr] 포함 구간
long long query(int ql, int qr) const {
long long sum = 0;
for (ql += n_, qr += n_; ql <= qr; ql /= 2, qr /= 2) {
if (ql % 2 == 1) sum += tree_[ql++]; // 왼쪽 경계가 오른쪽 자식이면 포함 후 한 칸 이동
if (qr % 2 == 0) sum += tree_[qr--]; // 오른쪽 경계가 왼쪽 자식이면 포함 후 한 칸 이동
}
return sum;
}
};이 레이아웃은 n이 2의 거듭제곱이 아니어도 합, 최솟값처럼 교환법칙이 성립하는 연산에서는 올바르게 동작합니다. 다만 행렬 곱처럼 순서가 중요한 연산에는 왼쪽과 오른쪽 결과를 따로 모아 합치는 처리가 추가로 필요합니다.
범용 세그먼트 트리 템플릿
병합 함수와 항등원만 바꾸면 합·최솟값·최댓값·GCD를 모두 같은 코드로 처리할 수 있습니다. 병합 함수와 항등원이 결합법칙을 만족하는 모노이드를 이뤄야 한다는 것이 유일한 조건입니다.
#include <algorithm>
#include <climits>
#include <vector>
template <typename T, typename Merge>
class SegmentTreeGeneric {
std::vector<T> tree_;
int n_;
Merge merge_;
T identity_;
void build(int node, int l, int r, const std::vector<T>& arr) {
if (l == r) { tree_[node] = arr[l]; return; }
int mid = (l + r) / 2;
build(2 * node, l, mid, arr);
build(2 * node + 1, mid + 1, r, arr);
tree_[node] = merge_(tree_[2 * node], tree_[2 * node + 1]);
}
T queryImpl(int node, int l, int r, int ql, int qr) const {
if (qr < l || r < ql) return identity_;
if (ql <= l && r <= qr) return tree_[node];
int mid = (l + r) / 2;
return merge_(queryImpl(2 * node, l, mid, ql, qr),
queryImpl(2 * node + 1, mid + 1, r, ql, qr));
}
void updateImpl(int node, int l, int r, int idx, const T& val) {
if (l == r) { tree_[node] = val; return; }
int mid = (l + r) / 2;
if (idx <= mid) updateImpl(2 * node, l, mid, idx, val);
else updateImpl(2 * node + 1, mid + 1, r, idx, val);
tree_[node] = merge_(tree_[2 * node], tree_[2 * node + 1]);
}
public:
SegmentTreeGeneric(const std::vector<T>& arr, Merge merge, T identity)
: tree_(4 * arr.size(), identity), n_(static_cast<int>(arr.size())),
merge_(merge), identity_(identity) {
if (n_ > 0) build(1, 0, n_ - 1, arr);
}
T query(int ql, int qr) const { return queryImpl(1, 0, n_ - 1, ql, qr); }
void update(int idx, const T& val) { updateImpl(1, 0, n_ - 1, idx, val); }
};
int main() {
std::vector<int> arr = {3, 1, 4, 1, 5, 9, 2, 6};
auto maxOp = [](int a, int b) { return std::max(a, b); };
SegmentTreeGeneric stMax(arr, maxOp, INT_MIN); // C++17 CTAD
// stMax.query(2, 5) == 9
}클래스 템플릿에서는 템플릿 인자 일부만 명시하고 나머지를 추론하게 할 수 없습니다. 그래서 SegmentTreeGeneric<int>(arr, lambda, INT_MIN)처럼 쓰면 컴파일 에러가 나고, 인자를 모두 생략해 C++17 클래스 템플릿 인자 추론(CTAD)에 맡기거나 SegmentTreeGeneric<int, decltype(maxOp)>처럼 둘 다 적어야 합니다.
응용: 인기순 자동완성과 주가 구간 분석
트라이로 접두사 후보를 모은 뒤 사용 빈도로 정렬하면 간단한 자동완성이 됩니다.
#include <algorithm>
#include <string>
#include <unordered_map>
#include <vector>
class AutocompleteService {
Trie trie_;
std::unordered_map<std::string, int> popularity_;
public:
void addWord(const std::string& word, int score = 1) {
if (trie_.insert(word)) popularity_[word] += score;
}
std::vector<std::string> suggest(const std::string& prefix, std::size_t limit = 10) const {
auto candidates = trie_.wordsWithPrefix(prefix);
auto byPopularity = [this](const std::string& a, const std::string& b) {
return popularity_.at(a) > popularity_.at(b);
};
if (candidates.size() > limit) {
std::partial_sort(candidates.begin(), candidates.begin() + limit,
candidates.end(), byPopularity);
candidates.resize(limit);
} else {
std::sort(candidates.begin(), candidates.end(), byPopularity);
}
return candidates;
}
};비교 함수에서 popularity_[a]를 쓰면 operator[]가 없는 키를 삽입할 수 있어 const 멤버 함수에서 쓸 수 없고, 실수로 맵을 키우기도 합니다. 트라이에 들어간 단어는 모두 맵에도 있으므로 at으로 조회했습니다. 상위 limit개만 필요하므로 전체를 정렬하는 대신 std::partial_sort를 썼는데, 후보가 k개일 때 O(k log limit)입니다. 다만 후보 수집 자체가 서브트리 크기에 비례하므로, 후보가 매우 많아지는 서비스에서는 노드에 상위 후보를 캐시하는 구조로 바꿔야 합니다.
두 자료구조를 함께 쓰는 예로, 구간 최저가는 세그먼트 트리에, 구간 거래량 합은 펜윅 트리에 맡길 수 있습니다.
class StockAnalyzer {
SegmentTreeMin minTree_;
FenwickTree volumeTree_;
public:
StockAnalyzer(const std::vector<int>& prices, const std::vector<long long>& volumes)
: minTree_(prices), volumeTree_(static_cast<int>(volumes.size())) {
volumeTree_.init(volumes);
}
int rangeMinPrice(int left, int right) const { return minTree_.query(left, right); }
long long rangeTotalVolume(int left, int right) const {
return volumeTree_.rangeSum(left + 1, right + 1); // 0-indexed → 1-indexed
}
void updatePrice(int index, int newPrice) { minTree_.update(index, newPrice); }
void addVolume(int index, long long delta) { volumeTree_.add(index + 1, delta); }
};바깥에서는 모두 0-indexed로 받고, 펜윅 트리를 호출하는 곳에서만 1을 더하도록 경계를 한 곳에 모은 것이 요점입니다. 인덱스 변환이 여러 곳에 흩어지면 앞에서 본 off-by-one 버그가 생기기 쉽습니다.
자주 묻는 질문 (FAQ)
Q. 구간 합에는 펜윅 트리와 세그먼트 트리 중 무엇을 써야 하나요?
A. 점 업데이트와 구간 합만 필요하다면 코드가 짧고 메모리를 적게 쓰는 펜윅 트리로 충분합니다. 펜윅 트리는 누적합의 차이로 구간 값을 구하기 때문에 구간 최솟값이나 최댓값처럼 역연산이 없는 쿼리에는 쓸 수 없으며, 이런 경우나 구간 업데이트가 필요한 경우에는 lazy propagation을 적용한 세그먼트 트리를 사용합니다. 어느 쪽이든 합이 int 범위를 넘을 수 있으면 long long으로 저장해야 합니다.
Q. 펜윅 트리에서 add(0, x)를 호출하면 왜 무한 루프에 빠지나요?
A. 펜윅 트리는 1-indexed가 표준이고, 갱신 루프가 i += (i & -i)로 다음 인덱스로 이동합니다. i가 0이면 i & -i도 0이라 인덱스가 전혀 증가하지 않아 루프가 끝나지 않습니다. 0-indexed 배열을 넣을 때는 ft.add(i + 1, arr[i])처럼 1을 더해 변환해야 합니다.
Q. 세그먼트 트리 배열은 왜 4n 크기로 잡나요?
A. 재귀로 구간을 반씩 나누는 세그먼트 트리는 n이 2의 거듭제곱이 아니면 한쪽이 한 단계 더 깊어져 노드 번호가 2n을 넘을 수 있습니다. 최대 번호는 n 이상인 가장 작은 2의 거듭제곱의 두 배보다 작으므로 4n이면 항상 충분합니다. 리프를 n..2n-1에 두는 반복문 구현은 2n으로 충분합니다.
참고 자료
- cp-algorithms: Segment Tree
- cp-algorithms: Fenwick Tree
- LeetCode 307. Range Sum Query - Mutable
- LeetCode 208. Implement Trie (Prefix Tree)