세그먼트 트리 (Java, Segment tree)

---

기업 코테에서 세그먼트 트리를 요구하는 경우는 흔치 않은 것 같지만

 

필요한 일이 생겨서

 

이참에 정리해보려고 한다..

 

값이 변하지 않는 데이터가 주어졌을 때,

 

구간 합을 빠르게 구하는 방법은   

 

prefix sum 을 이용하면 된다.

https://www.crocus.co.kr/843

구간 합(Prefix Sum) 알고리즘

 

만약, 데이터가 변한다면 Fenwick tree 를 구현해서 구간 쿼리를 효율적으로 해결할 수 있겠다.

https://yabmoons.tistory.com/438

[ 펜윅트리(FenwickTree) ] 개념과 구현방법 (C++)

 

하지만 구간에 대한 특정 값을 구하는 경우는 펜윅 트리를 이용해서 효율적으로 해결하기 어렵다.

 

예를 들어 구간 [A, B] 에서의 최댓값, 최솟값, gcd, lcm, 등등

 

이를 해결하기 위해서는 세그먼트 트리를 활용하는 것이 유리하다. (펜윅 트리보다 요구하는 메모리가 더 많긴 함)

 

구간 쿼리가 주어졌을 때, 

 

세그먼트 트리를 이용하면 

 

쿼리 연산과 업데이트 연산을 log(N) 에 해결할 수 있다.

 

⚠️ 하지만 Range update 에 대해서는 비효율적이다. M 개의 데이터가 있는 구간에 일괄적으로 업데이트를 수행한다고 하면 기존의 세그먼트 트리에서는 M*log(N) 이 필요하다. Range update 가 많이 발생하지 않는다면 큰 문제가 없겠지만, 만약 Range update 가 자주 발생하는 환경이라면 Lazy propagation 이라는 기법을 사용해야 효율적으로 해결할 수 있다. 본 포스팅에서는 그 경우까지는 고려하지 않는다.

 

설명은 아래에 주석으로 대체하도록 하겠다.

public static class segmentTree {
        long[] tree; // 트리의 노드 값이 저장될 배열
        int N; // 원본 데이터의 개수

        public segmentTree(int n, int[] arr) { // 생성자
            N = n;
            tree = new long[n * 4]; // 트리의 노드는 2N <= x <4N 이 보장된다. (top-down)
            build(arr, 1, 0, n - 1); // init
        }

        private long build(int[] arr, int node, int nodeLeft, int nodeRight) {
            if (nodeLeft == nodeRight) { // leaf 에 도달한 경우에 값을 넣는다.
                return tree[node] = arr[nodeLeft];
            }
            int mid = nodeLeft + (nodeRight - nodeLeft) / 2; // overflow 방지
            // (a+b)/2 라고 하면 (a+b) 에서 Long overflow 가 발생할 수 있다.
            long leftVal = build(arr, node * 2, nodeLeft, mid); 
            long rightVal = build(arr, node * 2 + 1, mid + 1, nodeRight);
            return tree[node] = merge(leftVal, rightVal);
        }

        public long merge(long left, long right) {
            return left + right; // sum
            // return Math.min(left, right); // min
            // return Math.max(left, right); // max
        }

        public Long query(int left, int right) {
            // 실제 입력 배열에서 [left, right] 구간에서의 query 결과값 반환
            return queryRecursive(left, right, 1, 0, N - 1);
        }

        public Long update(int index, int newValue) {
            return updateRecursive(index, newValue, 1, 0, N - 1); // 재귀적으로 업데이트
        }

        public Long updateRange(int left, int right, int newValue) { // range update
            return updateRangeRecursive(left, right, newValue, 1, 0, N - 1); 
        }

        private Long queryRecursive(int left, int right, int node, int nodeLeft, int nodeRight) {
            if (right < nodeLeft || nodeRight < left) {
                return 0L; // 해당 노드 미포함
            }
            if (left <= nodeLeft && nodeRight <= right) {
                return tree[node]; // 해당 노드 포함
            }
            int mid = nodeLeft + (nodeRight - nodeLeft) / 2;
            return merge(queryRecursive(left, right, node * 2, nodeLeft, mid),
                    queryRecursive(left, right, node * 2 + 1, mid + 1, nodeRight));
        }

        private Long updateRecursive(int index, int newValue, int node, int nodeLeft, int nodeRight) {
            if (index < nodeLeft || nodeRight < index) {
                return tree[node];
            }
            if (nodeLeft == nodeRight) { // 리프 노드에 도달하면
                return tree[node] = newValue;
            }
            int mid = nodeLeft + (nodeRight - nodeLeft) / 2;

            Long leftVal = updateRecursive(index, newValue, node * 2, nodeLeft, mid);
            Long rightVal = updateRecursive(index, newValue, node * 2 + 1, mid + 1, nodeRight);

            return tree[node] = merge(leftVal, rightVal);
        }

        private Long updateRangeRecursive(int left, int right, int newValue, int node, int nodeLeft, int nodeRight) {
            if (right < nodeLeft || nodeRight < left) {
                return tree[node];
            }
            if (nodeLeft == nodeRight) {
                return tree[node] = newValue;
            }
            int mid = nodeLeft + (nodeRight - nodeLeft) / 2;
            Long leftVal = updateRangeRecursive(left, right, newValue, node * 2, nodeLeft, mid);
            Long rightVal = updateRangeRecursive(left, right, newValue, node * 2 + 1, mid + 1, nodeRight);
            return tree[node] = merge(leftVal, rightVal);
        }
    }

이제 문제를 풀어보자.

 

위 코드만 숙지하면 간단하게 해결할 수 있다!

 

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

1275번: 커피숍2

세그먼트 트리는 자료구조 자체가 조금 난이도가 있기 때문에 기본 문제일지라도 골드 I 에서 시작한다.

 

( 백준 티어 올리기에 좋음 ,,,)

 

 

여튼 위 문제를 읽어보면,

 

문제 / 입력 설명

N 개의 원본 데이터가 주어지고 Q 개의 쿼리가 수행된다.

 

각각의 쿼리는 sum query 를 수행한 뒤 update 를 적용하는 것으로 정의된다.

 

위와 같이 입력이 주어지면, 데이터의 개수는 5개 이고 쿼리의 개수는 2개 인 것.

 

데이터는 [1, 2, 3, 4, 5] 이고

 

쿼리 1. 에서는 2~3 까지의 합을 구하고, 3 번쨰의 값을 1 으로 변경한다.

쿼리 2. 에서는 3~5 까지의 합을 구하고, 4 번째의 값을 1 으로 변경한다.

 

주의할 점은, 입출력에서 주어지는 순서 값들은 전부 1-indexed 이고

 

우리가 구현한 자료구조는 0-indexed 라는 것이다.

 

또한 구간 쿼리를 수행할 때, left 와 right 가 주어지면 left 가 right 보다 크게 주어지는 경우에는 left 와 right 를 바꾸어주어야 한다.

 

위에서 구현한 segmentTree 클래스를 이용하면 매우 쉽게 해결할 수 있다.

 

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer stk;
		// 백준 1275 커피숍 2 솔루션

        stk = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(stk.nextToken()); // 수의 개수
        int Q = Integer.parseInt(stk.nextToken()); // 턴의 개수
        int[] arr = new int[N];

        stk = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            arr[i] = Integer.parseInt(stk.nextToken());
        }

        segmentTree segmentTree = new segmentTree(N, arr);

        for (int i = 0; i < Q; i++) {
            stk = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(stk.nextToken());
            int y = Integer.parseInt(stk.nextToken());

            if (x > y) {
                int t = x;
                x = y;
                y = t;
            }

            int a = Integer.parseInt(stk.nextToken());
            int b = Integer.parseInt(stk.nextToken());

            Long sum = segmentTree.query(x - 1, y - 1);
            segmentTree.update(a - 1, b);
            bw.write(sum + "\n");
        }


        bw.flush();
        bw.close();
    }

 

 

 

 

---

참고 자료들

https://www.crocus.co.kr/648

세그먼트 트리(Segment Tree)

https://www.youtube.com/watch?v=075fcq7oCC8&t=938s