백준 21214번 Decoration

문제 요약

숫자 $N, K$가 주어진다. 다음을 만족하는 길이가 K인 수열 $s$중에서 모든 원소의 합이 가장 작은 $s$를 찾아서 출력하시오. 만족하는 수열이 없으면 -1을 출력하면 된다.

1. $s_i \neq s_j, i \neq j, 1 \le i,j \le K$
2. $ 0 \le s_i \lt N$
3. $\text{for all } 1 \le i \lt K, s_{i+1} \equiv s_i + D(s_i) \bmod N$, 여기서 $D(x)$는 $x$의 약수의 갯수를 말한다.

$ 1 \le N, K \le 1,000,000$

풀이

문제를 푸는 데에 있어서 중요한 관찰은 3번 조건때문에 $s_i$가 정해지면 $s_{i+1}$이 정해져 버린다는 점이다.

그래서 이런 관계를 함수 $f$로 표현해서 $f(x) = x + D(x) \bmod N$이라고 하자. $0 \le x \lt N$이다.

노드가 $N$개이고 각 노드에는 0부터 $N-1$까지 번호가 적혀있으며 $(x, f(x))$를 엣지로 갖는 functional graph를 생각하자.

$s_1$을 정한 뒤 수열 $s$를 만드는 과정은 $s_1$이 적혀있는 노드부터 시작해서 총 $K-1$번 엣지를 타고 넘어가면서 거쳐간 노드들의 번호를 $s_i$로 쓰는 것과 동일해진다.

그러면 이 수열의 합은 $s_1$부터 세서 간선을 $K-1$번 탈 때까지의 노드들의 번호를 다 합한 것이 된다.

0부터 N-1까지 자신의 노드부터 시작해서 이 누적합을 나이브하게 계산하면 $O(NK)$이지만 그래프가 functional graph이기 때문에 스파스 테이블을 이용하면 $O(NlogK)$만에 이것들이 계산 가능하다.

이제 1번 조건을 잘 봐야 한다. 우리가 만든 그래프는 사이클이 존재한다. 따라서, $K-1$번 간선을 타는 경로에서 등장하는 노드의 수가 정확히 $K$인지 확인해야 한다.

이는 DFS로 확인이 가능하다. 어떤 노드에서 DFS를 시작한 뒤에 사이클을 만나면 해당 사이클을 처리해주고 DFS를 끝내며 만나게 될 노드의 수를 계산할 수 있다.

시간복잡도를 계산해보자. $f$를 계산하는 데에는 $O(NlogN)$이 걸리며, 누적합을 계산하는 데에 $O(NlogK)$, 경로들에 $K$개의 노드가 등장하는 지 확인하는 데에 $O(N)$이 걸린다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MAXN = 1e6;
int N, K;
int trans[MAXN];
int cycle_num;
bool vis[MAXN];
ll max_len[MAXN];
ll cost[MAXN];
int cur_pos[MAXN], tmp_pos[MAXN], sp_table[MAXN];
ll sp_sum[MAXN], tmp_sum[MAXN];

void calc_trans() {
    for(int i=0;i<N;++i) trans[i] = i;
    for(int i=1;i<N;++i) {
        for(int j=i;j<N;j+=i) {
            ++trans[j];
        }
    }
    for(int i=0;i<N;++i) trans[i] %= N;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> K;
    calc_trans();
    for(int i=0;i<N;++i) cur_pos[i] = i;
    for(int i=0;i<N;++i) sp_table[i] = trans[i];
    for(int i=0;i<N;++i) sp_sum[i] = i;
    int k = 1;
    while(k <= K) {
        if(k & K) {
            for(int i=0;i<N;++i) {
                cost[i] += sp_sum[cur_pos[i]];
                cur_pos[i] = sp_table[cur_pos[i]];
            }
        }
        for(int i=0;i<N;++i) {
            tmp_pos[i] = sp_table[i];
            tmp_sum[i] = sp_sum[i];
        }
        for(int i=0;i<N;++i) {
            sp_table[i] = tmp_pos[tmp_pos[i]];
            sp_sum[i] += tmp_sum[tmp_pos[i]];
        }
        k <<= 1;
    }
    for(int i=0;i<N;++i) {
        if(!vis[i]) {
            vector<int> st;
            int cur = i;
            while(!vis[cur]) {
                vis[cur] = true;
                st.push_back(cur);
                cur = trans[cur];
            }
            if(max_len[cur] == 0) { // new cycle
                int cycle_start = -1;
                for(int i=0;i<st.size();++i) {
                    if(cur == st[i]) {
                        cycle_start = i;
                        break;
                    }
                }
                assert(cycle_start != -1);
                int cycle_len = st.size() - cycle_start;
                for(int i=cycle_start;i<st.size();++i) max_len[st[i]] = cycle_len;
            }
            while(!st.empty()) {
                int v = st.back();
                st.pop_back();
                if(max_len[v] == 0) {
                    assert(max_len[trans[v]] != 0);
                    max_len[v] = max_len[trans[v]]+1;
                }
            }
        }
    }
    pair<ll, int> ans = {1e18, -1};
    for(int i=0;i<N;++i) if(max_len[i] >= K) {
        ans = min(ans, {cost[i], i});
    }
    if(ans.second == -1) cout << -1 << '\n';
    else {
        int x = ans.second;
        for(int i=0;i<K;++i) {
            cout << x << " \n"[i==K-1];
            x = trans[x];
        }
    }
    return 0;   
}