[C++] 백준 2042번 - 구간 합 구하기

 

📖 문제

 

📋 코드

#include<iostream>
using namespace std;
typedef long long ll;
int N, M, K, a, b;
ll c;
ll input[1 << 20];
ll seg[1 << 21];

ll init(int node, int s, int e)
{
    if (s == e) return seg[node] = input[s]; // start 와 end 의 위치가 일치하면 input[start] 값을 넣어준다.
    int mid = (s + e) / 2;
    return seg[node] = init(2 * node, s, mid) + init(2 * node + 1, mid + 1, e); 
}

ll prefixSum(int node, int s, int e, int l, int r) // 시작노드, start, end, left, right
{
    if (e < l || r < s) return 0; // 찾아야하는 구간(l~r)과 현재 구간(s~e)이 포함되지 않을때
    if (l <= s && e <= r) return seg[node]; // 찾아야하는 구간(l~r)내에 현재 구간(s~e)이 포함될 때
    int mid = (s + e) / 2;
    // 찾아야하는 구간(l~r)이 현재 구간(s~e)에 포함되거나, 부분적으로 겹치는 경우
    return prefixSum(2 * node, s, mid, l, r) + prefixSum(2 * node + 1, mid + 1, e, l, r);
}

void changeNum(int node, int s, int e, int index, ll diff) {// 시작노드, start, end, 찾아야하는 곳(index), 변화량(diff)
    if (e < index || index < s) return; // 찾아야하는 곳(index)과 현재 구간(s~e)이 포함되지 않을때
    seg[node] += diff;
    if (s != e) {
        int mid = (s + e) / 2;
        changeNum(2 * node, s, mid, index, diff);
        changeNum(2 * node + 1, mid + 1, e, index, diff);
    }
    
}

int main()
{
    cin >> N >> M >> K;

    for (int i = 1; i <= N; i++) {
        cin >> input[i];
    }

    init(1, 1, N); // 세그먼트 트리 생성

    for (int i = 0; i < M + K; i++)
    {
        cin >> a >> b >> c;
        ll diff = 0;
        switch (a)
        {
        case 1: // 숫자 변경
            diff = c - input[b];
            input[b] = c;
            changeNum(1, 1, N, b, diff); // 시작노드, start, end, 찾아야하는 곳(index), 변화량(diff)
            break;
        case 2: // 구간합 출력
            cout << prefixSum(1, 1, N, b, c) << '\n';
            break;

        default:
            break;
        }
    }
}

👨🏻‍💻 결과

 

 

📕 풀이 방법

주의할 점!

int형을 범위를 넘어갈 수 있는 변수들은 long long으로 선언해서 자료형 범위를 초과하지 않도록 해줘야 합니다!

 

#include<iostream>
using namespace std;
typedef long long ll;
int N, M, K, a, b;
ll c;
ll input[1 << 20];
ll seg[1 << 21];

N의 개수가 최대 100만 개이므로 input배열 크기를 2^20개,
세그먼트 트리 값들을 담을 seq 배열 크기를 2^21개로 선언해줍니다.

 

ll init(int node, int s, int e)
{
    if (s == e) return seg[node] = input[s]; // start 와 end 의 위치가 일치하면 input[start] 값을 넣어준다.
    int mid = (s + e) / 2;
    return seg[node] = init(2 * node, s, mid) + init(2 * node + 1, mid + 1, e); 
}

init함수는 초기에 세그먼트 트리를 만들어주는 역할을 합니다.

// 예제 입력 1
5 2 2
1
2
3
4
5
1 3 6
2 2 5
1 5 2
2 3 5

예제 입력 1로 만들어진 세그먼트 트리

 

ll prefixSum(int node, int s, int e, int l, int r) // 시작노드, start, end, left, right
{
    if (e < l || r < s) return 0; // 찾아야하는 구간(l~r)과 현재 구간(s~e)이 포함되지 않을때
    if (l <= s && e <= r) return seg[node]; // 찾아야하는 구간(l~r)내에 현재 구간(s~e)이 포함될 때
    int mid = (s + e) / 2;
    // 찾아야하는 구간(l~r)이 현재 구간(s~e)에 포함되거나, 부분적으로 겹치는 경우
    return prefixSum(2 * node, s, mid, l, r) + prefixSum(2 * node + 1, mid + 1, e, l, r);
}

prefixSum 함수는 구간합을 구해주는 함수입니다. 주석을 읽어보시면 이해가 되실 겁니다! ㅎㅎ

 

void changeNum(int node, int s, int e, int index, ll diff) {// 시작노드, start, end, 찾아야하는 곳(index), 변화량(diff)
    if (e < index || index < s) return; // 찾아야하는 곳(index)과 현재 구간(s~e)이 포함되지 않을때
    seg[node] += diff;
    if (s != e) {
        int mid = (s + e) / 2;
        changeNum(2 * node, s, mid, index, diff);
        changeNum(2 * node + 1, mid + 1, e, index, diff);
    }
}

changeNum 함수는 세그먼트 트리를 업데이트해주는 함수입니다. 위 함수에서 핵심은 최상단의 노드에서 해당 인덱스를 가진 노드로 이동하면서 거쳐가는 모드 노드들에 변화를 주는 것입니다.

현재 2번째 값이 2인데 2를 6으로 변경하는 과정을 살펴보자면 2 -> 6으로 변경되므로 diff값은 +4가 되고 최상단 노드에서 해당 노드까지 이동하면서 거쳐가는 모든 노드 값을 모두 +4를 해주면서 업데이트를 해줍니다.

이러한 업데이트 과정의 연산 횟수는 최댓값이 노드의 최대 깊이(max 21)이므로 시간 복잡도가 logN이고 시간 초과를 피할 수 있습니다.

 

int main()
{
    cin >> N >> M >> K;

    for (int i = 1; i <= N; i++) {
        cin >> input[i];
    }

    init(1, 1, N); // 세그먼트 트리 생성

    for (int i = 0; i < M + K; i++)
    {
        cin >> a >> b >> c;
        ll diff = 0;
        switch (a)
        {
        case 1: // 숫자 변경
            diff = c - input[b];
            input[b] = c;
            changeNum(1, 1, N, b, diff); // 시작노드, start, end, 찾아야하는 곳(index), 변화량(diff)
            break;
        case 2: // 구간합 출력
            cout << prefixSum(1, 1, N, b, c) << '\n';
            break;

        default:
            break;
        }
    }
}

main함수에서 input값을 입력받아 세그먼트 트리를 생성하고 a, b, c를 입력받아 값 업데이트, 구간합 출력을 해줍니다.

 

 

🔗 링크

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

2042번: 구간 합 구하기