[PS] 백준 2042번 구간 합 구하기

고범수·2023년 3월 14일
0

Problem Solving

목록 보기
26/31

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

문제 설명


N(N <= 1000,000)개의 수가 순서대로 주어진다. 이 때 p ~ q번째 수 까지의 합을 구하는 연산과 x번째 수를 변경하는 연산이 있다. 이 떄, M + K (M <= 10,000, K <= 10,000)쿼리가 입력으로 들어오면 그 결과를 출력해야 하는 문제이다.

문제 풀이


순수하게 접근하면 아래와 같다.

  1. 주어진 구간에 속한 수의 합을 구하는데 O(N)
  2. 수를 변경하는 데 O(1)
  3. 전체 복잡도를 구하면 O(NK + M)

    시간 초과가 난다는 것을 알 수 있다.

누적합을 사용하면 아래와 같다.

  1. 주어진 구간에 속한 수의 합을 구하는데 O(1)
  2. 수를 변경하는 데 O(N)
  3. 전체 복잡도를 구하면 O(K + MN)

    수의 합을 구하는 연산의 복잡도 카테고리를 어떻게든 낮춰야 뭐라도 가능하다.

세그먼트 트리(Segment Tree)를 이용해야 한다.

무엇인지는 아래 사이트에서 학습하자.

https://book.acmicpc.net/ds/segment-tree

Segment Tree를 이용하면 아래와 같다.

  1. N개의 수의 합을 구하는데 O(logN)
  2. 수를 변경하는 데 O(logN)
  3. 전체 복잡도를 구하면 O(KlogN + MlogN)
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.StringTokenizer;

public class Main {

    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static StringTokenizer st;

    static class SegmentTree {
        long[] tree;

        public SegmentTree(int size) {
            this.tree = new long[size + 1];
        }

        long init(long[] input, int node, int start, int end) {
            if (start == end) {
                return tree[node] = input[start];
            }

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

        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 tree[node];
            }

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

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

            tree[node] += diff;

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

    static int n, m, k;
    static long[] input;

    public static void main(String[] args) throws Exception {
        st = new StringTokenizer(br.readLine());

        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());

        input = new long[n + 1];

        int size = (int) Math.pow(2, Math.ceil(log2(input.length)) + 1) - 1;
        SegmentTree segTree = new SegmentTree(size);

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

        segTree.init(input, 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());

            long sum = -1;
            if (a == 1) {
                long diff = c - input[b];
                input[b] = c;
                segTree.update(1, 1, n, b, diff);
            } else {
                sum = segTree.sum(1, 1, n, b, (int) c);
                bw.write(Long.toString(sum));
                bw.newLine();
            }
        }

        bw.flush();
    }

    static double log2(int x) {
        return Math.log(x) / Math.log(2);
    }
}

0개의 댓글