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

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

예제 입출력

• 예제 입력 1

5 2 2
1
2
3
4
5
1 3 6
2 2 5
1 5 2
2 3 5

• 예제 출력 1

17
12

Solution

#include <iostream>
#define MAX 1000001
using namespace std;

long long Arr[MAX * 4];

void update(int node, long long val){
    while(node > 0){
        Arr[node] += val;
        node /= 2;
    }
}

long long get(int s, int e){
    long long ans = 0;
    while(s <= e){
        if(s % 2 == 1) ans += Arr[s];
        if(e % 2 == 0) ans += Arr[e];
        s = (s + 1) / 2;
        e = (e - 1) / 2;
    }
    return ans;
}

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);

    int N, M, K;
    cin >> N >> M >> K;
    int s = 1;
    for(s ; s < N; s <<= 2);
    for(int i = 0; i < N; i++) cin >> Arr[i+s];
    for(int i = s - 1; i >= 1; i--) Arr[i] = Arr[2*i] + Arr[2*i+1];
    for(int i = 0; i < M + K; i++){
        int a;
        long long b, c;
        cin >> a >> b >> c;
        if(a == 1){
            update(b + s - 1, c - Arr[b + s - 1]);
        }
        else{
            cout << get(b + s - 1, c + s - 1) << '\n';
        }
    }
    return 0;
}

Segment Tree 대표 문제다. 솔직히말해서 겁나 어렵다. 애초에 Segment Tree라는 개념 자체가 그렇게 쉽지 않다. Index Tree를 활용하여 해당 문제를 풀어보았다.

특정 구간의 합을 구하는 방법엔 여러가지가 있다. Segment Tree든 Index Tree든 말이다. 아니면 그냥 선형적으로 떄려박는 것도 방법이다. 허나 그렇게 풀 수 있을 것 같으면 골드 문제가 아니다.

Segment Tree를 사용하지 않고 Index Tree를 활용한 이유는...사실 그게 더 코드가 짧아서? 뭣보다 내가 아직 Segment Tree를 아주 완벽하게 숙지하고 있는 것은 아니다. 코드로 구현하자니 아직까진 부족한 것 같아서, 시험 볼 일이 있어서 어떻게든 빠르게 외울 수 있는 Index Tree를 주로 공부한 것 또한 이유였다.

#include <iostream>
#define MAX 1000001
using namespace std;

long long Arr[MAX * 4];

Tree 역할을 해 줄 배열을 선언한다. 데이터의 범위가 64비트이므로 long long으로 선언한다. Leaf Node에 최소 N개의 데이터가 들어가야 하기 때문에 최대 N * 4를 한다.

void update(int node, long long val){
    while(node > 0){
        Arr[node] += val;
        node /= 2;
    }
}

Node 가 0보다 클 때 까지, 즉 root 가 될 때 까지 특정 값을 더한다. 즉 Parent들에 해당 위치에서의 업데이트 값을 갱신한다는 것이다. 특정 범위의 합을 더하는 Index Tree의 특성 상 반드시 해당 노드의 Parent까지 값을 업데이트 해 줘야 한다.

long long get(int s, int e){
    long long ans = 0;
    while(s <= e){
        if(s % 2 == 1) ans += Arr[s];
        if(e % 2 == 0) ans += Arr[e];
        s = (s + 1) / 2;
        e = (e - 1) / 2;
    }
    return ans;
}

s는 시작 인덱스, e는 끝 인덱스다. S와 E는 각각 특정 노드의 Left Child거나 Right Child일 수 있다. Left는 Parent 2의 인덱스 값을 가지고 Right 는 Parent 2 + 1의 인덱스 값을 가진다. Start는 Right index일 때만 값을 더해주고 End 는 Left index일 때만 값을 더해준다.

int s = 1;
for(s ; s < N; s <<= 2);

우선 트리의 Leaf Node에 데이터를 입력해야 한다. 즉 Leaf노드의 인덱스 값을 만들어주어야 한다. 그래서 N보다 작을 때 까지(?) s에 2를 제곱해서 인덱스 값을 만들어준다.

for(int i = 0; i < N; i++) cin >> Arr[i+s];
for(int i = s - 1; i >= 1; i--) Arr[i] = Arr[2*i] + Arr[2*i+1];

그 후 데이터를 입력받는다. 그리고 입력받은 값을 기반으로 인덱스 트리를 만들어준다. 간단하다. Parent 기준으로 인덱스를 하나씩 Root 까지 올려가며 그 Parent의 Left와 Right Child 값을 더해주면 된다. s-1부터니까 Leaf Node 이후로 순서대로 하나씩 쌓아올라간다고 생각하면 편하다.

for(int i = 0; i < M + K; i++){
    int a;
    long long b, c;
    cin >> a >> b >> c;
    if(a == 1){
        update(b + s - 1, c - Arr[b + s - 1]);
    }
    else{
        cout << get(b + s - 1, c + s - 1) << '\n';
    }
}

그 후 문제의 조건에 따라 M + K 만큼 데이터를 입력받는다. a는 명령어라 생각해서 굳이 int형에서 벗어날 필욘 없다 생각했다. 하지만 업데이트가 일어나는 데이터들은 조건에 따라 long long 형태를 띄어야 한다. 간단하다. 1이면 해당 데이터를 특정 위치의 값을 C로 업데이트. 그냥 무작정 값을 바꾸는게 아니라 변화량 만큼 변화를 줘야 하므로 저렇게 입력하였다. b또한 Leaf Node에서부터 업데이트 해 나가야 하므로 b + s - 1로 처리.
get또한 마찬가지. Leaf Node 기준이므로 b + s - 1 부터 c + s - 1 간에 값을 구한다. Bottom up 방식으로 데이터를 하나하나 찾아 올라가서 최종 구간 합을 구한다.

Index Tree를 연습 해 볼 수 있는 중요한 문제다. 반드시 풀어볼것.