[BOJ10999] 구간 합 구하기2 (lazy propagation)

문제 설명

문제 링크

• 구간 합 구하기 1과 같이 세그먼트 트리를 이용하여 푸는 문제이다.
• 다만 구간 합 구하기 1은 중간에 수 변경이 하나의 원소에 대해서 일어났다.
• 구간 합 구하기 2는 수 변경이 구간으로 주어진다.

풀이 과정

• 구간 합 구하기 1처럼, 수 변경이 하나의 원소에 대해서만 일어나면, 세그먼트 트리 역시 logN(트리의 높이)만큼만 내려가며 수정을 하면 된다.
• 하지만, 수 변경이 구간으로 주어지게 되면, 최악의 경우 O(NlogN)이 될 것이다.
• 쿼리마다 O(NlogN)을 사용하게 되고, 모든 업데이트 쿼리의 시간복잡도는 O(NMlogN)이 될 것이다.
• 이를 위해서 우리는 lazy propagation 알고리즘을 사용해야 한다.
• 간단히 말해서, 구간의 값이 업데이트가 되었을 때, 필요한 깊이만큼만 업데이트를 시키고, 나머지 자식노드의 값은 필요할 때 업데이트 하는 방식이다.
• 나중에 업데이트 해야되는 것은 트리 노드의 크기만큼 배열을 만들고, 그곳에 저장을 해둔다.
• lazy propagation의 상세한 설명은 이곳을 참조하자.
• lazy propagation의 구현은 매우 간단하다.

코드

import java.io.*;
import java.util.*;

public class BOJ10999_구간합구하기2 {
    static long[] input;
    static long[] tree;
    static long[] lazy;
    static int N, M, K;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder output = new StringBuilder();
        StringTokenizer 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];
        tree = new long[4 * N];
        lazy = new long[4 * N];
        for (int n = 1; n <= N; n++) {
            input[n] = Long.parseLong(br.readLine());
        }
        init(1, N, 1);
        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());
            int c = Integer.parseInt(st.nextToken());
            if (a == 1) {
                long d = Long.parseLong(st.nextToken());
                modify(1, N, b, c, d, 1);
            } else {
                output.append(sum(1, N, b, c, 1));
                output.append("\n");
            }
        }
        System.out.println(output);
    }

    static void lazyPropagate(int node, int start, int end) {
        if (lazy[node] == 0) return;
        tree[node] += (end - start + 1) * lazy[node];
        if (start != end) {
            lazy[node * 2] += lazy[node];
            lazy[node * 2 + 1] += lazy[node];
        }
        lazy[node] = 0;
    }
    static void modify(int start, int end, int left, int right, long diff, int node) {
        lazyPropagate(node, start, end);
        if (start > right || end < left) return;
        if (left <= start && end <= right) {
            tree[node] += (end - start + 1) * diff;
            if (start != end) {
                lazy[node * 2] += diff;
                lazy[node * 2 + 1] += diff;
            }
            return;
        }
        int mid = (start + end) / 2;
        modify(start, mid, left, right, diff, node * 2);
        modify(mid + 1, end, left, right, diff, node * 2 + 1);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }
    static long sum(int start, int end, int left, int right, int node) {
        lazyPropagate(node, start, end);
        if (start > right || end < left) {
            return 0;
        }
        if (left <= start && end <= right) {
            return tree[node];
        }
        int mid = (start + end) / 2;
        return sum(start, mid, left, right, node * 2) + sum(mid + 1, end, left, right, node * 2 + 1);
    }
    static long init(int left, int right, int node) {
        if (left==right){
            return  tree[node] = input[left];
        }
        int mid = (left + right) / 2;
        return tree[node] = init(left, mid, node * 2) + init(mid + 1, right, node * 2 + 1);
    }
}