[백준] 1275번. 커피숍2(JAVA)

문제

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

---

풀이(20분)

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

public class Main {

    static int N, Q, squareNum;
    static long[] segment;

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

        st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());
        int pow = 0;
        while(Math.pow(2, pow) < N){
            pow++;
        }
        segment = new long[(int)Math.pow(2,pow+1)];

        squareNum = (int)Math.pow(2, pow);

        st = new StringTokenizer(br.readLine());
        for(int i = squareNum; i < squareNum + N; i++){
            segment[i] = Long.parseLong(st.nextToken()); // 2^k 부터 값 저장.
        }

        for(int i = squareNum * 2 - 1; i >= 2; i--){
            segment[i/2] += segment[i];
        }

        StringBuilder sb = new StringBuilder();
        for(int i = 0; i < Q; i++){
            st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken());
            int y = Integer.parseInt(st.nextToken());
            if(x > y){
                int temp = x;
                x = y;
                y = temp;
            }
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            sb.append(query(x,y)).append("\n");
            update(a,b);
        }
        System.out.println(sb);
    }

    static long query(int start, int end){
        start += squareNum-1; // 리프 노드 시작 인덱스로 이동
        end += squareNum-1; // 리프 노드 끝 인덱스로 이동
        long sum = 0;

        while(start <= end){
            if(start % 2 == 1){
                sum += segment[start];
                start ++;
            }
            if(end % 2 == 0){
                sum += segment[end];
                end--;
            }
            start /= 2;
            end /= 2;
        }
        return sum;
    }

    static void update(int startIdx, int num){
        startIdx += squareNum - 1;
        long diff = num - segment[startIdx];

        while(startIdx > 0){
            segment[startIdx] += diff;
            startIdx /= 2;
        }
    }

}

---

문제 풀이 전략

 

구간합을 질의하고, 업데이트하기를 반복하는 대표적인 세그먼트 트리 문제이다.

💡 정확한 개념을 숙지하고 싶다면 아래 링크를 참고하자.

[도구정리] 세그먼트 트리(Segment Tree) 알고리즘

---

여기서는 세그먼트 트리에 대한 간략한 이해와 함께 문제를 해결해 보자.

{3, 1, 4, 1, 5, 9, 2, 6}이라는 배열이 존재할 때의 세그먼트 트리 초안이다.

우리는 주어진 배열을 이진트리로 만든 뒤 이를 이용하여 구간합을 탐색할 것이다.

이 과정에서 질의에 대한 시간복잡도는 비슷할지 몰라도 업데이트 시 O(logN)이라는 엄청난 성능을 보여준다.

 

자세한 개념은 건너뛸 것이기에 이진트리의 크기를 만드는 방법은 뒤로 하고, 완성본을 살펴보자.

인덱스 별로 하위 두 개의 합을 갖게 된다.

 

📖 업데이트 

우리는 이를 이용하여 업데이트를 하게 된다면 더욱 효율적일 거라는 게 보이는가?

예를 들어 원본 배열에서 값이 5인 인덱스 12를 업데이트한다면(다시 말하지만 세그먼트 트리의 인덱스에 대한 개념은 위의 링크를 참고하자.) 12, 6, 3, 1 4개 만을 업데이트하면 끝난다.

-> 물론 여기선 12, 13, 14, 15 인덱스를 업데이트하는 것과 동일할 수 있지만 배열의 크기가 커지거나 9번 정도의 인덱스라고 생각해 보자!

 

결론 : 세그먼트 트리는 구간합차곱의 경우 업데이트에 매우 유리하다.

---

📖 질의 방식

업데이트가 유용하다는 것은 알았으니 이제 질의 방식을 알아보자.

내가 원하는 구간합은 1~4번 인덱스의 합이다.(인덱스는 0부터 시작한다고 하자.)

 

startIdx = 1, endIdx = 4가 될 것이다. 여기서 +8을 각각 해주자.(인덱스를 맞추기 위함)

startIdx = 9, endIdx = 12가 된다.

 

💡규칙

1. startIdx <= endIdx일 때 반복한다.

2. startIdx -> 즉, 왼쪽에서 올라가는 것은 오른쪽 자식일 경우만 자기 자신의 값을 저장해 두고, startIdx++를 한다.

2-1. endIdx ->  즉, 오른쪽에서 올라가는 것은 왼쪽 자식일 경우만 자기 자신의 값을 저장해두고, endIdx--를 한다.

3. startIdx, endIdx를 모두 2로 나누어준다.

---

 

한 단계씩 반복하며 따라가 보자.

값을 저장해 두기 위해 sum을 쓰자.

 

1. startIdx = 9 -> 오른쪽 자식이므로 sum += segment[startIdx]를 해주자. sum == 1이다.

1-1. startIdx++ -> startIdx == 10이다.

2. endIdx = 12 -> 왼쪽 자식이므로 sum += segment[endIdx]를 해주자. sum == 6이다.

2-2. endIdx-- -> endIdx == 11이다.

3. startIdx /= 2, endIdx /=2로 startIdx == 5, endIdx == 5가 된다.

4. 아직 startIdx <= endIdx이므로 반복한다.

5. startIdx = 5 -> 오른쪽 자식이므로 sum += segment[startIdx]를 해주자. sum == 11이다.

5-1. startIdx++ -> startIdx == 6이다.

5. endIdx = 5 -> 오른쪽 자식이므로 아무 동작을 하지 않는다.

6. startIdx /= 2, endIdx /=2로 startIdx == 3, endIdx == 2가 된다.

7. 끝난다.

 

총합은 11이 된다.

질의 역시 깔끔하게 마무리된다.

 

그렇기에 구간합을 사용하는데 업데이트가 필요하다면 세그먼트 트리를 이용하자!!

---

이를 문제에 적용하기만 하면 된다.

 

질문 메서드와, 업데이트 메서드를 작성하여 문제를 해결하면 된다!!

자세한 구현은 코드를 참고하자.