[백준 16975] 수열과 쿼리 21

세그먼트 트리 응용문제

www.acmicpc.net/problem/16975

www.acmicpc.net/problem/2042

구간 합 세그먼트 트리를 처음 접한다면 위 문제를 먼저 푸는 것을 권장합니다.

---

[문제 풀이]

 

수열이 주어지면 최대 100,000개의 쿼리를 처리하는 문제입니다.

 

쿼리는 두 종류입니다.

  1. 특정 범위에 k를 더함

  2. 수열의 특정한 위치의 값 출력

 

1번 쿼리 때문에 느리게 갱신되는 세그먼트 트리 문제처럼 보입니다.

실제로 느리게 갱신되는 세그먼트 트리로 풀 수 있습니다.

 

하지만 2번 쿼리를 잘 보면 구간의 합을 구하는 것이 아닌 특정한 값을 찾는 문제입니다.

2번 쿼리마저 구간의 합을 구하는 문제였다면 이 문제는 빼도 박도 못하고 느리게 갱신되는 세그먼트 트리 문제였겠지만, 2번 쿼리가 조금 단순한 형태라 약간의 변형을 통해 일반적인 세그먼트 트리로 해결이 가능합니다.

 

 

세그먼트 트리, 느리게 갱신되는 세그먼트 트리는 이제 모두 잊고, 이 문제에 대해서 본질적으로 다가가 보겠습니다.

 

이 문제가 원하는 것은 1번 쿼리에 대한 변화를 기록하고 2번 쿼리에 대해 (i번째 수열의 원래 값) + (i번째 위치에 대한 변화량)을 출력하는 것입니다.

예를 들어 1, 2, 3, 4, 5라는 수열이 있을 때 [1, 5]에 3을 더하고 [3, 5]에 -2를 더한다면 수열의 두 번째 값은 2 + 3을 출력하고, 네 번째 값은 2 + (3 - 2)를 출력하면 답을 구할 수 있습니다.

 

그렇다면 남은 문제는 1번 쿼리에 대한 변화량을 어떻게 기록하느냐입니다.

변화량을 기록하는 것이므로 그래프로 표현하는 것이 적합해 보입니다.

 

 

1번 쿼리에 대한 변화 예시를 그래프로 표현해보겠습니다.

 

초기 상태

[1, 6]의 변화 초기 상태입니다.

아직은 아무것도 더하지 않아서 모든 범위에 대해 0입니다.

 

 

1 1 3 2 입력 후

[1, 3] 범위에 대해서 2를 더해주는 쿼리는 위 그림처럼 생각할 수 있습니다.

[1, 3]은 +2가 되었고 [4, 6]은 그대로입니다.

 

조금 생각을 바꾸면 x축 양의 방향으로 평행 이동하는 선을 x = 1에서 잠깐 2만큼 올리고,  x = 4가 될 때 다시 2만큼 내린 것처럼 볼 수 있습니다.

우선은 이런 생각을 할 수가 있다 정도만 하고 넘어갑시다.

 

 

1 2 5 1 입력 후

위의 상태에서 [2, 5] 범위에 1을 더하는 쿼리는 위와 같습니다.

1과 6을 제외한 모든 부분이 1이 더해졌습니다.

 

 

우선 쿼리가 여기서 끝났다고 생각하고 그래프를 분석해봅시다.

 

이 그래프의 성질은 아까 위에서 생각한 내용처럼 선이 x축 양의 방향으로 계속 이동하는 것처럼 생각할 수 있다는 것입니다.

그렇기 때문에 i번째에 k를 더하면 i의 y축(변화량)만 올라가지 않고 i + 1 이후의 모든 변화량에 대해서도 k가 더해진 상태가 되는 것이고, 우리는 1번 쿼리에서 주어진 범위에 맞게  y축(변화량)을 범위 시작지점에 +k, 범위 끝지점에 -k를 더해주는 것만으로도 변화량 그래프를 그릴 수 있던 것입니다.

 

두 번째 쿼리로 예시를 든다면 [2, 5]에 1을 더하는데 2, 3, 4, 5에 모두 1을 더하지 않고 2에서 1을 더하고 6에서 1을 빼주는 것만으로도 그래프가 그려졌습니다.

그 이유는 그래프의 선이 양의 방향으로 계속 평행이동을 했기 때문이고,

바꿔 말해서 이전의 y축 변화가 이후의 y축에도 그대로 영향을 미치기 때문입니다.

 

슬슬 뭔가 떠오르지 않을까요?

 

놀랍게도, 세그먼트 트리로 우리가 방금 했던 그래프를 그리는 과정을 그대로 구현할 수 있습니다.

1. [a, b]에 k를 더하는 것은 세그먼트 트리의 a번째 원소에 +k, b + 1번째 원소에 -k

2. 수열의 i번째 원소의 변화량을 구하는 것은 0부터 i번째까지의 구간합을 세그먼트 트리로 구하는 것

 

즉, 느리게 갱신하는 세그먼트 트리가 아닌 일반적인 세그먼트 트리로 약간의 변형을 주어 위 문제를 해결할 수 있습니다.

 

---

[소스 코드] / Java

 

 

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

public class Main {
    private static int[] arr;
    private static long[] sumtree;

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

        arr = new int[N + 1];
        sumtree = new long[4 * N + 1];

        StringTokenizer st = new StringTokenizer(br.readLine(), " ");
        for (int i = 1; i < N + 1; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        int Q = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();
        while (Q-- > 0) {
            st = new StringTokenizer(br.readLine(), " ");
            char q = st.nextToken().charAt(0);
            if (q == '1') {
                int u = Integer.parseInt(st.nextToken()), v = Integer.parseInt(st.nextToken()), k = Integer.parseInt(st.nextToken());
                update(1, N, 1, u, k);
                update(1, N, 1, v + 1, -k);
            } else {
                int v = Integer.parseInt(st.nextToken());
                long result = getSum(1, N, 1, 1, v) + arr[v];
                sb.append(result).append('\n');
            }
        }

        System.out.println(sb.toString());
        br.close();
    }

    private static long getSum(int left, int right, int node, int u, int v) {
        if (v < left || right < u) return 0;
        else if (u <= left && right <= v) return sumtree[node];
        int mid = (left + right) / 2;
        return getSum(left, mid, node * 2, u, v) + getSum(mid + 1, right, node * 2 + 1, u, v);  
    }

    private static long update(int left, int right, int node, int i, int k) {
        if (left <= i && i <= right) {
            if (left == right) return sumtree[node] += k;
            int mid = (left + right) / 2;
            sumtree[node] = update(left, mid, node * 2, i, k) + update(mid + 1, right, node * 2 + 1, i, k);
        }
        return sumtree[node];
    }
}