5299. 아름다운 균일 수

수 N이 있고 이 수를 k진법으로 표현했을 때 모든 자리의 수가 동일하고 2자리 이상일 때 N을 k의 균일수 라고 정의한다.

f(N)은 N을 균일수로 만드는 모든 k값의 합을 뜻하며 문제에서 묻는것은

 이다.

N의 제한이 1e9이므로 1부터 1e9까지 모두 구하면 당연하게도 시간초과가 난다.

1. 아이디어

하나씩 다 구하는건 앞서 말한 것 처럼 TLE가 뜨기 때문에 다른 방법을 생각해야 한다.

균일수의 성질에 대해서 잘 생각해보자. k진법일 때 1로 이루어진 길이가 L인 수는

k^(L-1) + k^(L-2) + k^(L-3) + k^(L-4) + ....  + k + 1

이다. 만약 이게 1로 이루어진게 아니라 2로 이루어진 거라면?

2 * (k^(L-1) + k^(L-2) + k^(L-3) + k^(L-4) + ....  + k + 1)

가 된다. 저건 k-1로 이루어진 형태까지도 변하지 않는 패턴이다. 이 패턴에서 아이디어를 하나 얻을 수 있다.

라는 배열을 만들어서 arr[i][j] = i진법에서 길이가 j인 1로 이루어진 수를 미리 구해놓는다.

그 다음에는 수 N이 주어졌을 때 

1. 모든 자리가 a로 이루어져 있고 

2. 길이가 n일 때 

3. upperbound ( N/a )를 이분탐색으로 찾는다.

이렇게 하면 시간복잡도 만에 

길이가 n이고 모든 자리가 a인 균일수가 N보다 작거나 같다는 조건 하에 몇진법까지 표현 가능한가를 말해준다.

예를들어 N = 100, a = 1, n = 3일 때 이분탐색을 할 배열 모양새는

{ 7, 13, 21, 31, 43, 57, 73, 91, 111,,,, } 이고

{ 111(2), 111(3), 111(4), 111(5), 111(6), 111(7), 111(8), 111(9), 111(10),,,, }

을 뜻한다. 그러므로 100보다 작거나 같고 3자리면서 111 형태인 균일수는 총 9개이며

이는 a+1 ~ upperbound ( N/a ) 까지의 등차수열 합과 같다.

이대로 짜보자. TLE가 떴다. 그 뒤에 틀렸습니다가 뜬 것도 있으니 시간복잡도상 문제가 되는건지는 잘 모르겠다.

어쨋든 TLE가 떴으니 시간을 줄일만한걸 생각해 보자.

이분탐색을 할 때 2자리 짜리 균일수도 구했는데 이게 생각보다 시간을 많이 먹는 것 같다.

미리 배열을 만들어 놓을 때 [][2], [][3]만 31622까지고 나머지는 엄청 작으니 이걸 줄이면 제법 빨라질 것 같다.

2자리수는 가만보면 단순한 식으로도 계산이 가능했다. 모든 수 N은 N-1진법으로 11(N-1)로 표현이 가능하기 때문이다.

#include <iostream>
using namespace std;
 
typedef unsigned long long ull;
 
ull tc;
ull mul[32000][28];
ull idxs[28];
 
void init() {
    for (register ull i = 2; i*i <= 1e9; ++i) {
        ull cnt = 0, pd = i * i*i, temp = i + 1 + i * i;
        while (temp <= 1e9) {
            mul[i][cnt] = temp;
            idxs[cnt] = i;
            cnt++;
            temp += pd;
            pd *= i;
        }
    }
}
 
inline ull run(ull N) {
    register ull ans = 0;
    for (register ull a = 1; a*a < N; ++a) {
        if (a + 1 >(N - a) / a) continue;
        ans += (((N - a) / a)*(((N - a) / a) + 1)) / 2 - (a * (a + 1)) / 2;
        for (register ull n = 0; n < 27; ++n) {
            if (mul[a + 1][n] * a > N) break;
            ull l = a + 1, r = idxs[n], mid, cc = 0;
            while (l <= r) {
                mid = (l + r) / 2;
                if (mul[mid][n] * a <= N) {
                    cc = mid;
                    l = mid + 1;
                }
                else
                    r = mid - 1;
            }
            if (cc > 1)
                ans += (cc * (cc + 1)) / 2 - (a * (a + 1)) / 2;
        }
    }
    return ans;
}
 
int main() {
    ios::sync_with_stdio(false); cin.tie(0);
 
    init();
    cin >> tc;
 
    for (ull t = 1; t <= tc; ++t) {
        ull N, result;
        cin >> N;
        result = run(N);
        cout << "#" << t << " " << result << '\n';
    }
 
    return 0;
}