[백준/BOJ C++] 17633번 : 제곱수의 합 (More Huge) (정수론, 폴라드 로, Diamond IV)

이 포스트는 프로그래밍 문제 사이트 백준 Online Judge(BOJ)의 17633번 : '제곱수의 합' 문제의 풀이 코드와 해설을 다루고 있습니다.

 

문제 난이도는 Solved.ac 기준 Diamond IV에 해당하며, 문제를 풀이하기 위해 각종 제곱수 관련 정리와 폴라드 로 인수분해 알고리즘에 대한 이해가 필요합니다.

 

 

17633번 : 제곱수의 합 (More Huge)

17633번: 제곱수의 합 (More Huge)

 

10^18 이하의 자연수가 입력되었을 때, 이 자연수를 최소 몇 개의 제곱수들의 합으로 나타낼 수 있는지를 출력하는 문제입니다.

불과 한 줄로 표현이 가능한 짧은 문제이지만, 어떻게 풀이의 방향을 잡아야할지 막막합니다.

 

저의 경우에는 비슷한 문제인 '제곱수의 합 2 (More Huge)' 문제의 난이도가 Ruby IV로 책정된 것을 보고, 제곱수들을 직접 구하는 것보다 개수만을 구하는 것이 훨씬 쉽다는 사실을 유추할 수 있었습니다.

그러한 생각을 바탕으로, 제곱수의 개수를 구하는 과정은 단순히 어떤 정리를 활용해서 풀이할 수 있기 때문이라고 유추할 수 있었습니다.

그래서 관련된 정리들을 찾아보았고 관련 식들에 대해 짧게나마 공부해서 풀이할 수 있었습니다.

 

문제를 자세히 보면 아주 중요한 힌트를 주었는데, 바로 라그랑주가 모든 자연수는 4개 이하의 제곱수의 합으로 표현할 수 있다고 한 정리입니다.

사실 라그랑주의 네 제곱수의 합 정리와 같이, 역사적인 증명들을 즉석에서 증명해가면서 문제를 풀이하는 것은 불가능합니다.

따라서 이 문제에서는 제곱수의 합과 관련된 Square Theorem들을 간단히 활용하여 문제를 풀이해볼 것입니다.

 

 

 

문제를 풀이하는 데 필요한 정리들만 가져와서 한 장에 깔끔하게 정리했습니다.

사실 위의 정리들의 증명 과정을 이해하기 위해서는 아주 긴 시간이 소요됩니다.

따라서 여기서는 그 과정을 생략하도록 하고 간단히 설명만 하도록 하겠습니다.

위에서 사용된 Fermat's two-sqaure theorem, Legendre's three-square theorem, Lagrange's four-square theorem에 대해서는 각각 검색해보면 더 많은 정보를 얻을 수 있습니다.

 

문제를 풀이하기 전에 각 정리를 사용하기 위해서는 N에 대한 소인수분해가 필요합니다.

따라서 List라는 벡터를 선언해주고, 여기에 아주 큰 수를 작은 시간에 소인수분해 가능한 폴라드 로 인수분해 알고리즘을 사용하여 저장해준 뒤 오름차순으로 정렬을 수행해줍니다.

 

N이 1개의 제곱수의 합으로 표현되기 위해서는, 그냥 N이 제곱수이면 됩니다.

다만 C++에서 지원하는 sqrt 함수의 경우 int 범위만 다루기 때문에, 소인수분해를 활용하여 제곱수인지를 검사해주어야 합니다.

만약 소인수를 같은 2개씩 짝지을 수 있으면, 이들을 둘로 동일하게 나누어 곱한 값들이 제곱근이 되므로 N이 1개의 제곱수의 합으로 표현될 수 있음을 알 수 있습니다.

 

N이 2개의 제곱수의 합으로 표현할 수 있는지 알기 위해서는 Fermat's two-sqaure theorem을 사용해야 합니다.

Fermat's two-sqaure theorem이란 4k+3꼴의 소인수의 지수가 모두 짝수이면 N을 2개의 제곱수의 합으로 표현할 수 있다는 정리입니다.

따라서 각 소인수를 4로 나눈 나머지가 3인 소인수가 짝수개인지를 확인해주면 됩니다.

 

N이 3개의 제곱수의 합으로 표현할 수 있는지 확인하기 위해서는 Legendre's three-square theorem을 사용해야 합니다.

Legendre's three-square theorem이란 N이 4^x(8k+7)꼴이 아니면 N을 3개의 제곱수의 합으로 표현할 수 있다는 정리입니다.

따라서 N을 4로 나눌 수 있을 때까지 나눠준 뒤 8로 나눈 나머지가 7이 아닌지를 확인해주면 됩니다.

 

N이 4개의 제곱수의 합으로 표현할 수 있는지 확인하기 위해서는 Lagrange's four-square theorem을 사용해야 합니다.

이는 문제에서도 주어진 정리로, 모든 자연수는 4개 이하의 자연수의 제곱들의 합으로 표현할 수 있음이 보장되기 때문에, 위의 3가지 경우에 해댱되지 않으면 4개의 자연수의 제곱의 합으로 표현됨을 알 수 있습니다.

 

 

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

ull Power(__int128 x, __int128 y, __int128 mod) { // ret = (x^y)%mod
    x %= mod;
    __int128 ret = 1;
    while(y > 0) {
        if(y%2 == 1) ret = (ret*x)%mod;
        x = (x*x)%mod;
        y /= 2;
    }
    return (ull)ret;
}

bool checkPrime(ull n, ull a) {
    if(a%n == 0) return true;
    ull k = n-1;
    while(1) {
        ull temp = Power(a, k, n);
        if(temp == n-1) return true;
        if(k%2 == 1) return (temp == 1 || temp == n-1);
        k /= 2;
    }
}

bool isPrime(ull n) {
    if(n == 2 || n == 3) return true;
    if(n%2 == 0) return false;

    ull a[] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37};
    for(int i=0; i<12; i++)
        if(!checkPrime(n, a[i])) {
            return false;
            break;
        }
    return true;
}

ull GCD(ull a, ull b) {
    if(a < b) swap(a, b);
    while(b != 0) {
        ull r = a%b;
        a = b;
        b = r;
    }
    return a;
}

ull findDiv(__int128 n) {
    if(n%2 == 0) return 2;
    if(isPrime(n)) return n;

    __int128 x = rand()%(n-2) + 2, y = x, c = rand()%10 + 1, g = 1;
    while(g == 1) {
        x = (x*x%n + c)%n;
        y = (y*y%n + c)%n;
        y = (y*y%n + c)%n;
        g = GCD(abs(x-y), n);
        if(g == n) return findDiv(n);
    }
    if(isPrime(g)) return g;
    else return findDiv(g);
}

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

    ull N, temp;
    cin >> N; temp = N;

    vector<ull> List;
    while(N > 1) {
        ull div = findDiv(N);
        List.push_back(div);
        N /= div;
    }
    sort(List.begin(), List.end());

    bool check = true;
    if(List.size()%2 == 0) {
        for(int i=0; i<List.size(); i+=2)
            if(List[i] != List[i+1]) check = false;
        if(check) { cout << "1"; return 0; }
    }

    check = true; int exp_cnt = 1;
    for(int i=1; i<List.size(); i++) {
        if(List[i] == List[i-1]) exp_cnt++;
        else {
            if(List[i-1]%4 == 3 && exp_cnt%2 != 0) check = false;
            exp_cnt = 1;
        }
    }
    if(List[List.size()-1]%4 == 3 && exp_cnt%2 != 0) check = false;
    if(check) { cout << "2"; return 0; }

    N = temp;
    while(N%4 == 0) N /= 4;
    if(N%8 != 7) cout << "3";
    else cout << "4";
}

 

남은 과정은, 코드로써 위의 정리를 분기별로 나누어 처리해주면 됩니다.

다만 10^18과 같이 큰 자연수를 소인수분해하기 위해 폴라드 로 알고리즘이 사용되기 때문에, 이를 선행적으로 구현해야 합니다.

 

 

 

풀이를 제출해보면 24ms의 시간이 소요되고 모든 테스트케이스를 통과함을 확인할 수 있습니다.