[자바]백준 2042번 구간 합 구하기[세그먼트 트리][엄탱]

안녕하세요. 개발자 엄탱입니다.
이 글은 알고리즘을 공부하면서 공부 기록용입니다.
그래서 설명마다 일기용으로 편하게 작성하여 반말 형식으로 작성하려고 합니다.
그리고 보시다가 더 좋은 방법이나 잘 못 알고 있는 내용이 있다면 알려주시면 정말 감사하겠습니다.
좋은 하루 되세요 :)

문제 링크

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

2042번: 구간 합 구하기

문제 설명

구간 합 구하기 문제는 단순하게 문제에서 주어지는 구간의 합을 산출하는 문제이다.

단, 중간 중간 수의 변경이 일어나고 그에 맞게 새로운 합을 산출하면 된다.

해설

그냥 단순히 구간 합을 구하겠다 생각하고 구간이 구해질 때마다 반복문을 돌려서 답을 산출하겠다 생각하고 반복문을 돌리면 수의 개수가 최대 2^63 이기 때문에 시간초과가 날 수 있다.

이 문제는 전형적인 세그먼트 트리 문제이다.

우선, 세그먼트 트리를 공부하고 연습용으로 풀면 좋은 문제이다.

 

아래는 백준에서 세그먼트 트리에 대해서 알려주는 글이다.

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

세그먼트 트리 (Segment Tree)

코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine(), " ");

        int N = Integer.parseInt(st.nextToken()); // 수의 개수
        int M = Integer.parseInt(st.nextToken()); // 수의 변경이 일어나는 횟수
        int K = Integer.parseInt(st.nextToken()); // 구간의 합을 구하는 횟수

        long[] arr = new long[N + 1];

        for (int i = 1; i <= N; i++) {
            arr[i] = Long.parseLong(br.readLine());
        }

        SegmentTree segment = new SegmentTree(arr);
        segment.init(1, 1, N);

        for (int i = 0; i < M + K; i++) {
            st = new StringTokenizer(br.readLine(), " ");
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long c = Long.parseLong(st.nextToken());

            if (a == 1) {
                segment.update(1, 1, N, b, c - segment.getArr()[b]);
                segment.updateArr(b, c);
            } else {
                System.out.println(segment.sum(1, 1, N, b, (int)c));
            }
        }

        br.close();
    }
}

class SegmentTree {
    private long[] tree;
    private long[] arr;

    public SegmentTree(long[] arr) {
        int height = (int)Math.ceil(Math.log(arr.length - 1) / Math.log(2));

        tree = new long[(int)Math.pow(2, height + 1)];
        this.arr = arr;
    }

    public long init(int node, int start, int end) {
        if (start == end) {
            return this.tree[node] = this.arr[start];
        }

        return this.tree[node] = init(node * 2, start, (start + end) / 2) +
                init(node * 2 + 1, (start + end) / 2 + 1, end);
    }

    public void update(int node, int start, int end, int idx, long diff) {
        if (idx < start || end < idx) {
            return;
        }

        this.tree[node] += diff;

        if (start != end) {
            update(node * 2, start, (start + end) / 2, idx, diff);
            update(node * 2 + 1, (start + end) / 2 + 1, end, idx, diff);
        }
    }

    public long sum(int node, int start, int end, int left, int right) {
        if (left > end || right < start) {
            return 0;
        }

        if (left <= start && end <= right) {
            return this.tree[node];
        }

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

    public long[] getArr() {
        return this.arr;
    }

    public void updateArr(int idx, long value) {
        this.arr[idx] = value;
    }
}