C++ 제곱근 분할법 알고리즘 (Sqaure-root Decomposition, 예제 풀이 포함)

제곱근 분할법 알고리즘(Square-root Decomposition)이란 구간을 원소의 제곱근의 수로 나눠서 특정 구간의 데이터를 다루거나 연산, 탐색하는 방법의 알고리즘입니다.

 

 

 

sqrt(N)개의 연속한 원소를 가지는 bucket 여러 개를 가지고 쿼리들을 처리해주는 것입니다.

예를 들어 원소의 수가 16개일 때, 하나의 bucket의 크기를 sqrt(16) = 4로 잡으면 1~4, 5~8, 9~12, 13~16 구간의 원소를 가지는 4개의 bucket을 잡을 수 있습니다.

 

만약 구간 합을 구하는 구하는 문제를 푼다고 생각해보면, 각 bucket이 가지는 값은 그 bucket이 가지는 원소들의 합으로 둘 수 있습니다.

예를 들어 arr[i] = i (1 <= i <= 16)이라면, bucket[1] = 1 + 2 + 3 + 4, bucket[2] = 5 + 6 + 7 + 8, ... 이런 식입니다.

 

원소를 갱신하려면 원소와 bucket의 값을 모두 갱신해주어야 합니다.

예를 들어 arr[2] = 5로 갱신해봅시다.

그러면 arr[2]가 포함되어있는 bucket[1]에 대해서도 값을 갱신해주어야 하므로, bucket[1] = 0으로 초기화 해준 뒤, arr[1] ~ arr[4]의 값을 다시 더해줍니다.

 

 

 

이제 구간합을 구하는 방법을 생각해봅시다.

쿼리를 통해 4 ~ 13번째 원소들의 합을 구한다고 가정합시다.

그러면 4번째 원소와 13번째 원소는 일일이 더해주고, 가운데 bucket의 구간이 통째로 포함된 구간은 미리 구해둔 각 bucket의 구간 합을 더해서 구해주는 방식입니다.

그러면 arr[4] + ... + arr[15] = arr[4] + bucket[2] + bucket[3] + arr[13]으로 구할 수 있습니다.

 

이렇게 하면 하나의 쿼리에 대해 O(√N) 시간에 처리할 수 있습니다.

세그먼트 트리를 이용하면 O(log N) 시간에 쿼리를 처리할 수 있으므로, N이 커질수록 세그먼트 트리에 비해 불리해지지만 적절한 N에 대해서는 충분히 통과할 수 있습니다.

참고로 하나의 bucket의 수를 반드시 sqrt(N)개로 설정할 필요는 없고 시간 제한에 따라 시간 복잡도 계산을 통해 bucket의 크기를 임의로 설정해줄 수도 있습니다.

 

자세한 예시를 드는 것이 좋겠지만 그림 그리기는 귀찮으니 예제 코드를 가지고 정리해보겠습니다.

** 아직 알고리즘의 활용에 익숙하지 않습니다만, 나중에 까먹으면 혼자 복습하려고 정리해봅니다.

 

 

백준 BOJ 2042번 : 구간 합 구하기

위에서 예시로 든 가장 간단한 문제인 구간 합을 구하는 쿼리들을 처리하는 문제를 풀이해봅시다.

 

다음은 1 a b 쿼리에 대해서는 a번째 원소를 b로 갱신 해주고, 2 a b에 대해서는 a번째 원소부터 b번째 원소까지의 구간 합을 구하여 출력해주는 프로그램의 코드입니다.

 

 

#include <bits/stdc++.h>
#define int long long
using namespace std;

vector<int> v, u;

main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL), cout.tie(NULL);

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

    v.resize(N);
    for(int i=0; i<N; i++) cin >> v[i];

    int S = sqrt(N); u.resize(N/S + 1);
    for(int i=0; i<N; i++) u[i/S] += v[i];

    while(M--) {
        int Q; cin >> Q;

        if(Q == 1) {
            int idx, val; cin >> idx >> val;
            idx--;

            v[idx] = val;

            u[idx/S] = 0;
            int l = (idx / S) * S, r = l + S;
            for(int i=l; i<r; i++) u[idx/S] += v[i];
        }
        else if(Q == 2) {
            int l, r; cin >> l >> r;
            l--, r--;

            int sum = 0;

            while(l % S != 0 && l <= r) sum += v[l++];
            while((r + 1) % S != 0 && l <= r) sum += v[r--];

            while(l <= r) {
                sum += u[l/S];
                l += S;
            }

            cout << sum << "\n";
        }
    }
}

 

먼저 bucket의 크기 S를 sqrt(N)의 내림값으로 잡아주고, bucket들에 해당하는 2차원 벡터 u의 크기를 N/S + 1로 잡아줍니다. (N/S로 설정하면 크기를 벗어날 수 있습니다.)

각 bucket에 대해 해당 bucket이 담당하는 구간의 합을 구하여 저장해줍니다.

 

Q = 1, 즉 원소를 갱신하는 쿼리에 대해서는, 원소는 그냥 v[idx] = val로 갱신해주면되고, bucket의 경우에는 위에서 언급했듯 해당 bucket을 0으로 초기화하고 bucket에 해당하는 구간의 합을 다시 구해서 저장해주어야 합니다.

 

Q = 2, 구간 합을 구하는 쿼리에 대해서는 먼저 bucket의 일부를 차지하는 구간부터 처리해줍니다.

왼쪽 구간에 대해서는 l % S = 0이 될 때까지 l을 증가시키며 단일 원소의 값을 더해주고, 오른쪽 구간에 대해서는 (r + 1) % S = 0이 될 때까지 r을 감소시키며 단일 원소의 값을 더해줍니다.

이렇게 l과 r이 bucket의 경계까지 이동하게 되면 이제 bucket에 이미 저장되어있는 구간 합들만 bucket 단위로 더해주면 됩니다.

이렇게 계산할 경우 모든 원소의 합을 구하는 경우에도 루트 N개의 bucket에만 접근하게 되므로 O(√N) 시간에 처리가 가능합니다.

 

 

세그먼트 트리로 처리했을 때와의 수행 시간을 비교해보겠습니다.

아래가 세그먼트 트리로 풀이한 경우, 위가 제곱근 분할법으로 풀이한 경우의 시간입니다.

단일 쿼리를 로그 시간에 처리하는 세그먼트 트리가 비교적 빠름을 알 수 있습니다.

 

다만 소스 코드의 길이를 비교해보면 제곱근 분할법으로 구현하는 것이 조금 더 간단할 수 있다는 장점이 있습니다.

 

 

+ 제곱근 분할법을 사용하여 풀이할 수 있는 문제들 중에서 다음과 같은 쿼리 문제들 역시 풀이해보았으니, 더 어려운 난이도의 예제 풀이가 필요하다면 참고해보시는 것도 좋을 것 같습니다.

 

백준 수열과 쿼리 18 (BOJ 14504), 수열과 쿼리 1.5 (BOJ 17410) 풀이 (제곱근 분할법)