実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 234

問題概要

以下の条件を満たす長さ $N$ の整数列を辞書順で小さい方から全て出力せよ。

条件 :
- $i$ 番目の要素は $1$ 以上 $R_i$ 以下。
- 総和が $K$ の倍数である。

制約

入力はすべて整数。
$1 \leq N \leq 8$
$2 \leq K \leq 10$
$1 \leq R_i \leq 5$

考察

考えられる整数列は最大でも $5^8 = 390625 \simeq 3 \times 10^5$ なので、候補を全列挙する方針で考えてみる。

辞書順で小さい順に出力するので、初期状態が $\{1, 1, \cdots, 1\}$ という配列を作り、右端の要素から順に1ずつ加算していく。具体的なアルゴリズムは以下の通り。

配列 $\mathrm{ans}$ を $\{1, 1, \cdots, 1\}$ で初期化し、和 $\mathrm{sum}$ を $N$ とする。
配列のインデックスになるまで以下の操作を繰り返す。
- $\mathrm{sum}$ が $K$ の倍数ならば $\mathrm{ans}$ の中身を出力する。
- $\mathrm{idx} = N - 1$ とする。
- $\mathrm{ans}_{\mathrm{idx}} = R_{\mathrm{idx}}$ (上限に到達している) ならば、 $\mathrm{ans}_{\mathrm{idx}} = 1$ に戻し、 $\mathrm{idx}$ を $1$ 減らす (1個左を見る) 。 $\mathrm{sum}$ も調整。
- $\mathrm{idx} \lt 0$ を判定。
- $\mathrm{ans}_{\mathrm{idx}}$ を $1$ 増やし、 $\mathrm{sum}$ も $1$ 増やす。

コード

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

int main()
{
    int N, K;
    cin >> N >> K;
    vector<int> R(N);
    rep(i, 0, N) cin >> R[i];

    vector<int> ans(N, 1);
    int sum = N;
    while (true)
    {
        if (sum % K == 0)
        {
            for (int x : ans)
                cout << x << " ";
            cout << endl;
        }

        int idx = N - 1;
        while (idx >= 0 && ans[idx] == R[idx])
        {
            sum -= ans[idx];
            ans[idx] = 1;
            sum++;
            idx--;
        }

        if (idx < 0)
            break;

        ans[idx]++;
        sum++;
    }
}