2042번: 구간 합 구하기

2042번: 구간 합 구하기

사용 언어: C++

 

풀이

리프 노드들은 입력 값들, 부모 노드들은 값들의 합으로 설정해서 O(log n)으로 구간의 합을 구하는 세그먼트 트리 문제다. 주의할 점은 a == 1일때 c가 int 범위를 넘어서 오버플로우가 발생할 수 있어서 long long을 설정해야한다.

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;

// SZ는 N보다 크거나 같은 2^k 꼴의 수 (2^20 == 1,048,576 > 1,000,000)
constexpr int SZ = 1 << 20;
ll T[SZ << 1];

// x번째 수를 v로 지정
void Set(int x, ll v)
{
    x += SZ;
    T[x] = v;
    while (x /= 2)
    {
        T[x] = T[x * 2] + T[x * 2 + 1];  // 부모 노드 갱신
    }
}

// [l, r] 구간의 합을 구하는 함수
ll Sum(int l, int r)
{
    ll res = 0;
    for (l += SZ, r += SZ; l <= r; l /= 2, r /= 2)
    {
        if (l % 2 == 1) res += T[l++];  // l이 오른쪽 자식일 경우 해당 값을 더하고 다음으로 넘어감
        if (r % 2 == 0) res += T[r--];  // r이 왼쪽 자식일 경우 해당 값을 더하고 이전으로 넘어감
    }
    return res;
}

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

    int N, M, K;
    cin >> N >> M >> K;

    for (int i = 0; i < N; i++)
    {
        ll num;
        cin >> num;
        Set(i, num);  // 세그먼트 트리 초기화
    }

    for (int i = 0; i < M + K; i++)
    {
        int a, b;
        ll c;

        cin >> a >> b >> c;

        if (a == 1)
        {
            // b번째 수를 c로 변경
            Set(b - 1, c);
        }
        else if (a == 2)
        {
            // b번째부터 c번째까지의 합을 출력
            cout << Sum(b - 1, (int)c - 1) << '\n';
        }
    }

    return 0;
}