[백준] 2042번 구간 합 구하기 (JAVA)

문제

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

2042번: 구간 합 구하기

 

설명

1. 중간에 수의 변경이 빈번히 일어난다.
2. 그 중간에 어떤 부분의 합을 구하려 한다.

N개의 수가 있고 위의 두 조건을 M번 실행한다고 가정한다.

 

매번 연산을 새로 하는 일반적인 방법으로는 1번 조건의 수행시간은 O(1) 이고 2번 조건의 수행시간은 \(O(N)\) 이다. 최종적으로 \(O(NM)\) 의 시간복잡도를 가진다.

 

매번 합을 구하는 상황을 없애기 위해 누적합 배열을 이용하면 1번 조건의 수행시간은 O(N) 이고 2번 조건의 수행시간은 \(O(1)\) 이다. 최종적으로 \(O(NM)\) 의 시간복잡도를 가진다.

 

하지만 인덱스 트리를 이용하면 1번 조건의 수행시간은 \(O(logN)\) 이고 2번 조건의 수행시간은 \(O(logN)\) 이다. 최종적으로 \(O(MlogN)\) 의 시간복잡도를 가진다.

---

인덱스 트리는 Top-Down 또는 Bottom-Up 방식으로 구현할 수 있는데 Top-Down 방식으로 구현했다.

리프 노드(start == end)일 경우 해당 노드에 배열의 값을 저장하고 그 값을 리턴한다.

내부 노드일 경우에는 왼쪽 자식(node * 2)과 오른쪽 자식(node * 2 + 1) 값을 합쳐서 노드에 저장하고 그 값을 리턴한다.

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

    int mid = (start + end) / 2;
    // 왼쪽 자식(2 * node) 합 + 오른쪽 자식(2 * node + 1) 합
    return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
}

각 줄마다 입력된 첫 숫자에 따라 1이면 숫자 바꾸는 연산을 하고, 2이면 구간 합을 구하는 연산을 한다.

if (flag == 1) {
    int value1 = Integer.parseInt(st.nextToken());
    long value2 = Long.parseLong(st.nextToken());
    long diff = value2 - nums[value1];
    nums[value1] = value2;
    update(1, N, 1, value1, diff);
}
else {
    int value1 = Integer.parseInt(st.nextToken());
    int value2 = Integer.parseInt(st.nextToken());
    System.out.println(query(1, N, 1, value1, value2));
}

노드가 queryLeft와 queryRight 범위 밖에 있다면 연관이 없기 때문에 0을 리턴하고 범위 안에 쏙 들어있으면 현재 노드 값을 리턴한다. 범위에 걸쳐 있다면 왼쪽 query와 오른쪽 query를 호출하고 두 값을 합쳐서 리턴한다.

long query(int start, int end, int node, int queryLeft, int queryRight) {
    // 1. 연관 없음
    if (queryRight < start || end < queryLeft) {
        return 0;
    }

    // 2. 판단 가능 (쏙 들어있음)
    if (queryLeft <= start && end <= queryRight) {
        return tree[node];
    }

    // 3. 판단 불가 (걸쳐 있음)
    int mid = (start + end) / 2;
    return query(start, mid, node * 2, queryLeft, queryRight) + query(mid + 1, end, node * 2 + 1, queryLeft, queryRight);
}

노드에 target이 포함되지 않다면 바로 리턴하여 무시하고 포함된다면 현재 노드에 diff를 반영한다.

자식이 있을 경우 왼쪽 자식과 오른쪽 자식 모두 update를 호출하여 diff를 반영한다.

void update(int start, int end, int node, int target, long diff) {
    // 1. 연관 없음
    if (target < start || end < target)
        return;

    // 2. 연관 있음
    tree[node] += diff;

    // 아직 탐색할 수 있는 범위가 더 있다면 하위에 있는 노드도 바꿔줌
    if (start != end) {
        int mid = (start + end) / 2;
        update(start, mid, node * 2, target, diff);
        update(mid + 1, end, node * 2 + 1, target, diff);
    }
}

 

전체 코드

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

public class Main {

    static int N, M, K;
    static long[] nums;
    static long[] tree;

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

        int mid = (start + end) / 2;
        // 왼쪽 자식(2 * node) 합 + 오른쪽 자식(2 * node + 1) 합
        return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
    }

    static long query(int start, int end, int node, int queryLeft, int queryRight) {
        // 1. 연관 없음
        if (queryRight < start || end < queryLeft) {
            return 0;
        }

        // 2. 판단 가능 (쏙 들어있음)
        if (queryLeft <= start && end <= queryRight) {
            return tree[node];
        }

        // 3. 판단 불가 (걸쳐 있음)
        int mid = (start + end) / 2;
        return query(start, mid, node * 2, queryLeft, queryRight) + query(mid + 1, end, node * 2 + 1, queryLeft, queryRight);
    }

    static void update(int start, int end, int node, int target, long diff) {
        // 1. 연관 없음
        if (target < start || end < target)
            return;

        // 2. 연관 있음
        tree[node] += diff;

        // 아직 탐색할 수 있는 범위가 더 있다면 하위에 있는 노드도 바꿔줌
        if (start != end) {
            int mid = (start + end) / 2;
            update(start, mid, node * 2, target, diff);
            update(mid + 1, end, node * 2 + 1, target, diff);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

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

        tree = new long[N * 4];

        init(1, N, 1);

        for (int i = 0; i < M + K; i++) {
            st = new StringTokenizer(br.readLine());
            int flag = Integer.parseInt(st.nextToken());

            if (flag == 1) {
                int value1 = Integer.parseInt(st.nextToken());
                long value2 = Long.parseLong(st.nextToken());
                long diff = value2 - nums[value1];
                nums[value1] = value2;
                update(1, N, 1, value1, diff);
            }
            else {
                int value1 = Integer.parseInt(st.nextToken());
                int value2 = Integer.parseInt(st.nextToken());
                System.out.println(query(1, N, 1, value1, value2));
            }
        }
    }
}