[JAVA] BOJ 백준 2042번 - 구간 합 구하기

목차

• 문제
• 내가 푼 방법
• 자바 코드
• 결과 및 회고

---

1. 문제

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

2042번: 구간 합 구하기

---

2. 내가 푼 방법

다른 분들의 풀이를 보고 푼 문제이다. 구간에 맞는 값을 구할 때 세그먼트 트리에 어떻게 데이터가 저장되는지 이해했다.

 

먼저 트리를 세팅하는 initTree() 함수는 리프 노드가 채워지는 상태를 base condition으로 두고 재귀 형식으로 구현했다.

public static long initTree(int s, int e, int node) {
    if (s == e) return tree[node] = arr[s];

    int mid = (s + e) / 2;
    return tree[node] = initTree(s, mid, 2*node) + initTree(mid+1, e, 2*node+1);
}

 

 

다음은 인덱스 값을 업데이트하는 updateTree() 함수를 선언해 주었다. 인덱스 값이 구간을 넘어설 경우에는 업데이트되지 않으면서, 리프 노드까지 변경된 상태를 base condition으로 두었다.

public static void updateTree(int s, int e, int node, int idx, long diff) {
    if (idx < s || e < idx) return;

    tree[node] += diff;

    if (s == e) return;

    int mid = (s + e) / 2;
    updateTree(s, mid, 2*node, idx, diff);
    updateTree(mid+1, e, 2*node+1, idx, diff);
}

 

 

마지막으로 구간 값을 구하는 sumTree() 함수는 앞서 구간 값을 미리 구해두었으므로, 구하고자 하는 a ~ b 사이의 인덱스 범위에 속하게 된다면 바로 구간 값을 리턴해주는 것이다.

public static long sumTree(int s, int e, int node, int a, int b) {
    if (s > b || e < a) return 0;
    if (a <= s && e <= b) return tree[node];

    int mid = (s + e) / 2;
    return sumTree(s, mid, 2*node, a, b) + sumTree(mid+1, e, 2*node+1, a, b);
}

 

---

3. 자바 코드

깃허브 풀이 주소 : https://github.com/geujeog/BOJ/blob/main/B2042.java

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

class Main {
    static long[] arr;
    static long[] tree;

    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 st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken()) + Integer.parseInt(st.nextToken());
        arr = new long[N];
        tree = new long[4*N];

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

        initTree(0, N-1, 1);

        for (int k = 0; k < K; k++) {
            st = new StringTokenizer(br.readLine());

            int command = Integer.parseInt(st.nextToken());

            if (command == 1) {
                int idx = Integer.parseInt(st.nextToken()) - 1;
                long target = Long.parseLong(st.nextToken());

                updateTree(0, N-1, 1, idx, target - arr[idx]);
                arr[idx] = target;
            }
            else if (command == 2) {
                int a = Integer.parseInt(st.nextToken()) - 1;
                int b = Integer.parseInt(st.nextToken()) - 1;
                bw.write(sumTree(0, N-1, 1, a, b)+"\n");
            }
        }

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

    public static long sumTree(int s, int e, int node, int a, int b) {
        if (s > b || e < a) return 0;
        if (a <= s && e <= b) return tree[node];

        int mid = (s + e) / 2;
        return sumTree(s, mid, 2*node, a, b) + sumTree(mid+1, e, 2*node+1, a, b);
    }

    public static void updateTree(int s, int e, int node, int idx, long diff) {
        if (idx < s || e < idx) return;

        tree[node] += diff;

        if (s == e) return;

        int mid = (s + e) / 2;
        updateTree(s, mid, 2*node, idx, diff);
        updateTree(mid+1, e, 2*node+1, idx, diff);
    }

    public static long initTree(int s, int e, int node) {
        if (s == e) return tree[node] = arr[s];

        int mid = (s + e) / 2;
        return tree[node] = initTree(s, mid, 2*node) + initTree(mid+1, e, 2*node+1);
    }
}

---

4. 결과 및 회고

처음에 세그먼트 트리를 이용해 구간 값을 구하는 거 보고 진짜.. 어떻게 이런 생각을 했지라고 생각했는데ㅎㅎ

나도 더 풀다 보면 바로 알고리즘을 생각해 낼 수 있겠지..?