[C++] 세그먼트 트리(Segment Tree)

세그먼트 트리(Segment Tree)

배열 내 범위 값을 구하거나 특정 범위의 값들을 변경할 때 빠르게 접근하기 위해 사용한다. 이름 그대로 구간들을 보존하고 있는 트리이다.

• 리프 노드는 배열의 그 수 자체를 의미하고 리프 노드 외 노드들은 자식 노드를 대표하는 값이다.
• 리프 노드를 제외한 노드는 비워져있는데, 세그먼트 트리의 목적에 따라 채워야한다.

 

arr = [5, 2, 3, 6, 7, 8, 1, 11]

 

기본 원리 : 합

 

https://www.acmicpc.net/blog/view/9

세그먼트 트리 (Segment Tree)

 

1) 초기화 : 세그먼트 트리 만들기

// a: 배열 a
// tree: 세그먼트 트리
// node: 세그먼트 트리 노드 번호
// node가 담당하는 합의 범위가 start ~ end
long long init(vector<long long>& a, vector<long long>& tree, int node, int start, int end)
{
    // 리프 노드일 경우
    if (start == end)
        return tree[node] = a[start];

    int mid = (start + end) / 2;
    return tree[node] = init(a, tree, node * 2, start, mid) + init(a, tree, node * 2 + 1, mid + 1, end);
}

 

2) 쿼리 : 구간 검색

// node가 담당하는 구간이 start ~ end 이고, 구해야하는 합의 범위는 left ~ right 이다.
long long query(vector<long long>& tree, int node, int start, int end, int left, int right)
{
    // 1. 자식 노드 둘 다 범위 밖일 경우
    if (left > end || right < start)
        return 0;

    // 2. 자식 노드 둘 다 범위 안에 있을 경우   
    if (left <= start && end <= right)
        return tree[node];

    // 3. 자식 노드 중 하나만 범위 안에 있을 경우
    int mid = (start + end) / 2;
    return query(tree, node * 2, start, mid, left, right) + query(tree, node * 2 + 1, mid + 1, end, left, right);
}

 

3) 수 변경하기

long long update(vector<long long>& tree, int node, int start, int end, int index, long long value)
{
    // 1. start ~ end에 index가 포함되지 않는 경우
    if (index < start || index > end)
        return tree[node];

    // 2. start ~ end에 index가 포함되는 경우 중 리프 노드일 경우
    if (start == end)
        return tree[node] += value;

    // 3. start ~ end에 index가 포함되는 경우 중 리프 노드가 아닐 경우
    int mid = (start + end) / 2;
    return tree[node] = update(tree, node * 2, start, mid, index, value) + update(tree, node * 2 + 1, mid + 1, end, index, value);
}

 

문제 : 합

 

https://www.acmicpc.net/problem/3653

3653번: 영화 수집

#include <iostream>
#include <vector>
using namespace std;

#define SIZE 200000 // n+m

int n, m;
vector<int> tree, idx;

int init(int node, int start, int end, int left, int right)
{
    if (left > end || right < start)
        return 0;
    if (start == end)
        return tree[node] = 1;
    int mid = (start + end) / 2;
    return tree[node] = init(node * 2, start, mid, left, right) + init(node * 2 + 1, mid + 1, end, left, right);
}

int query(int node, int start, int end, int left, int right)
{
    if (left > end || right < start)
        return 0;
    if (left <= start && end <= right)
        return tree[node];
    int mid = (start + end) / 2;
    return query(node * 2, start, mid, left, right) + query(node * 2 + 1, mid + 1, end, left, right);
}

int update(int node, int start, int end, int index, int value)
{
    if (index > end || index < start)
        return tree[node];
    if (start == end)
        return tree[node] += value;
    int mid = (start + end) / 2;
    return tree[node] = update(node * 2, start, mid, index, value) + update(node * 2 + 1, mid + 1, end, index, value);
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int t;
    cin >> t;
    while (t--) {
        tree.clear();
        tree.resize(4 * SIZE + 1);
        idx.clear();
        idx.resize(SIZE / 2 + 1);

        cin >> n >> m;
        for (int i = 1; i <= n; i++)
            idx[i] = m + i;

        init(1, 1, m + n, m + 1, m + n);

        for (int i = 0; i < m; i++) {
            int num;
            cin >> num;
            cout << query(1, 1, m + n, 1, idx[num] - 1) << " ";

            update(1, 1, m + n, idx[num], -1); // 해당 디비디를 꺼내고 업데이트
            idx[num] = m - i;
            update(1, 1, m + n, idx[num], 1); // 해당 디비디를 맨 위에 넣고 업데이트
        }

        cout << '\n';
    }
}

 

문제 : 곱

 

https://www.acmicpc.net/problem/5676

5676번: 음주 코딩

#include <iostream>
using namespace std;

int tree[400004], arr[100001];

int init(int node, int start, int end)
{
	if (start == end)
		return tree[node] = arr[start];
	int mid = (start + end) / 2;
	return tree[node] = init(node * 2, start, mid) * init(node * 2 + 1, mid + 1, end);
}

int query(int node, int start, int end, int left, int right)
{
	if (left > end || right < start)
		return 1;
	if (left <= start && end <= right)
		return tree[node];
	int mid = (start + end) / 2;
	return query(node * 2, start, mid, left, right) * query(node * 2 + 1, mid + 1, end, left, right);
}

int update(int node, int start, int end, int index, int diff)
{
	if (index < start || index > end)
		return tree[node];
	if (start == end)
		return tree[node] = diff;
	int mid = (start + end) / 2;
	return tree[node] = update(node * 2, start, mid, index, diff) * update(node * 2 + 1, mid + 1, end, index, diff);
}

int main()
{
	cin.tie(NULL);
	cout.tie(NULL);
	ios::sync_with_stdio(false);

	int N, K;
	while (cin >> N >> K) {
		for (int i = 1; i <= N; i++) {
			cin >> arr[i];
			if (arr[i] > 0)
				arr[i] = 1;
			else if (arr[i] < 0)
				arr[i] = -1;
		}

		init(1, 1, N);

		while (K--) {
			char c;
			int i, j;
			cin >> c >> i >> j;

			if (c == 'C') {
				if (j > 0)
					j = 1;
				else if (j < 0)
					j = -1;

				update(1, 1, N, i, j);
				arr[i] = j;
			}
			else if (c == 'P') {
				int answer = query(1, 1, N, i, j);
				if (answer == 0)
					cout << 0;
				else if (answer > 0)
					cout << '+';
				else
					cout << '-';
			}
		}
		cout << '\n';
	}
}

 

문제 : 최솟값

 

https://www.acmicpc.net/problem/2336

2336번: 굉장한 학생

#include <iostream>
#include <algorithm>
using namespace std;

int scores[3][500001], tree[500001 * 4];
pair<int, int> firstTest[500001];

int query(int node, int start, int end, int left, int right)
{
    if (left > end || right < start)
        return 1e9;
    if (left <= start && end <= right)
        return tree[node];
    int mid = (start + end) / 2;
    return min(query(node * 2, start, mid, left, right), query(node * 2 + 1, mid + 1, end, left, right));
}

int update(int node, int start, int end, int index, int value) 
{
    if (index < start || index > end)
        return tree[node];
    if (start == end)
        return tree[node] = value;
    int mid = (start + end) / 2;
    return tree[node] = min(update(node * 2, start, mid, index, value), update(node * 2 + 1, mid + 1, end, index, value));
}

int main()
{
    cin.tie(NULL);
    cout.tie(NULL);
    ios::sync_with_stdio(false);

    int n;
    cin >> n;
    for (int i = 0; i < 3; i++) {
        for (int j = 1; j <= n; j++) {
            int num;
            cin >> num;
            scores[i][num] = j;
        }
    }

    for (int i = 1; i <= n * 4; i++)
        tree[i] = 1e9;

    // 첫번째 시험 등수대로 정렬
    for (int i = 1; i <= n; i++)
        firstTest[i] = { scores[0][i], i };
    sort(firstTest + 1, firstTest + n + 1);

    int ans = 0;
    for (int i = 1; i <= n; i++) {
        int index = firstTest[i].second;

        // 최솟값이 현재 업데이트 중인 학생의 세번째 시험 등수보다 크다면 굉장한 학생이다.
        if (query(1, 1, n, 1, scores[1][index]) > scores[2][index])
            ans++;
        update(1, 1, n, scores[1][index], scores[2][index]);
    }
    cout << ans;
}

 

문제 : 레이지 프로퍼게이션(Lazy Propagation)

 

https://www.acmicpc.net/problem/10999

10999번: 구간 합 구하기 2

#include <iostream>
#include <algorithm>
using namespace std;

const int SIZE = 1 << 21;

struct SegTree {
    int treeStart;
    long long tree[SIZE], lazy[SIZE];

    SegTree() 
    {
        treeStart = SIZE / 2;
        fill(tree, tree + SIZE, 0);
        fill(lazy, lazy + SIZE, 0);
    }

    void init()
    {
        for (int i = treeStart - 1; i > 0; i--)
            tree[i] = tree[i * 2] + tree[i * 2 + 1];
    }

    // 구간([start, end)) node의 lazy 값을 propagate
    void propagate(int node, int start, int end) 
    {
        if (lazy[node] != 0) {
            // 리프 노드가 아니면 자식들에게 lazy 미룸
            if (node < treeStart) {
                lazy[node * 2] += lazy[node];
                lazy[node * 2 + 1] += lazy[node];
            }

            tree[node] += lazy[node] * (end - start);
            lazy[node] = 0;
        }
    }

    // 구간([left, right))에 value를 더하라
    void update(int node, int start, int end, int left, int right, int value)
    {
        propagate(node, start, end);

        if (left >= end || right <= start)
            return;
        if (left <= start && end <= right) {
            lazy[node] += value;
            propagate(node, start, end);
            return;
        }

        int mid = (start + end) / 2;
        update(node * 2, start, mid, left, right, value);
        update(node * 2 + 1, mid, end, left, right, value);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    // 구간([left, right))의 합을 구하라
    long long sum(int node, int start, int end, int left, int right)
    {
        propagate(node, start, end);

        if (left >= end || right <= start)
            return 0;
        if (left <= start && end <= right)
            return tree[node];

        int mid = (start + end) / 2;
        return sum(node * 2, start, mid, left, right) + sum(node * 2 + 1, mid, end, left, right);
    }
} seg;

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int N, M, K;
    cin >> N >> M >> K;

    for (int i = 0; i < N; i++)
        cin >> seg.tree[seg.treeStart + i];
    seg.init();

    for (int i = 0, a, b, c, d; i < M + K; i++) {
        cin >> a;
        if (a == 1) {
            cin >> b >> c >> d;
            seg.update(1, 0, seg.treeStart, b - 1, c, d);
        }
        else {
            cin >> b >> c;
            cout << seg.sum(1, 0, seg.treeStart, b - 1, c) << '\n';
        }
    }
}

 

문제 : 머지 소트 트리

 

https://www.acmicpc.net/problem/13544

13544번: 수열과 쿼리 3

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

const int SIZE = 1 << 18;

struct MergeSortTree {
    vector<int> arr[SIZE];

    void construct()
    {
        for (int i = SIZE / 2 - 1; i > 0; --i) {
            vector<int>& c = arr[i], & l = arr[i * 2], & r = arr[i * 2 + 1];
            arr[i].resize(l.size() + r.size());

            for (int j = 0, p = 0, q = 0; j < c.size(); ++j)
                c[j] = (q == r.size() || (p < l.size() && l[p] < r[q])) ? l[p++] : r[q++];
        }
    }

    int greater(int left, int right, int value, int node = 1, int start = 0, int end = SIZE / 2)
    {
        if (end <= left || right <= start)
            return 0;
        if (left <= start && end <= right)
            return arr[node].end() - upper_bound(arr[node].begin(), arr[node].end(), value); // value 보다 큰 개수

        int mid = (start + end) / 2;
        return greater(left, right, value, node * 2, start, mid) + greater(left, right, value, node * 2 + 1, mid, end);
    }
}mst;

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int N;
    cin >> N;
    for (int i = 0; i < N; ++i) {
        int a;
        cin >> a;
        mst.arr[SIZE / 2 + i].push_back(a);
    }

    mst.construct();

    int M, L = 0;
    cin >> M;
    while (M--) {
        int i, j, k;
        cin >> i >> j >> k;
        L = mst.greater((i ^ L) - 1, j ^ L, k ^ L);
        cout << L << '\n';
    }
}