けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 367 C - Enumerate Sequences (4Q, 灰色, 300 点)

全探索:再帰関数 全探索 数列 全探索:再帰関数(ビット全探索困難) そのまま覚えたい典型問題 灰色diff AtCoder AtCoder300点 ABC-C NoviSteps4Q

全探索の基礎になるような再帰関数の使い方！

問題概要

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

  • i 番目の要素は 1 以上 R_{i} 以下
  • 総和が K の倍数

制約

  • 1 \le N \le 8
  • 2 \le K \le 10
  • 1 \le R_{i} \le 5

考えたこと

このように、長さ N のさまざまな数列を生成する方法については、次の記事に詳しく書いた！

drken1215.hatenablog.com

今回も、ここで書いた方法が使える。まず、 i 番目の要素が 1 以上 R_{i} 以下であるような数列を列挙するのは、次のような再帰関数で実現できる。なお、R は何らかの方法で与えられるものとする（下のコードでは、再帰関数の引数にしている）。

void dfs(vector<int> &A) {
    // 終端条件 --- N 回までいったら終わり
    if (A.size() == R.size()) {
        for (int i = 0; i < A.size(); ++i) cout << A[i] << " ";
        cout << endl;
        return;
    }

    for (int v = 1; v <= R[A.size()]; ++v) {
        A.push_back(v);
        dfs(A, R);
        A.pop_back(); // これが結構ポイント
    }
}

ここまでできたら、あとは終端条件を処理する部分において、「総和が K の倍数のときに限り出力する」というよううにすればよい。

計算時間を見積もろう。最悪ケースは次の通りである。

  • N = 8
  • K = 1
  • R = \lbrack 5, 5, \dots, 5 \rbrack

このとき、出力する数列の個数は 5^{8} 個となる。この程度の分量ならば余裕で間に合う。

コード

まず、入力として与えられる K, R は再帰関数の引数とした。さらに、 A の途中経過における総和 sum も引数として与えることとした。

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

void dfs(vector<int> &A, int K, const vector<int> &R, int sum) {
    // 終端条件 --- N 回までいったら終わり
    if (A.size() == R.size()) {
        if (sum % K == 0) {
            for (int i = 0; i < A.size(); ++i) cout << A[i] << " ";
            cout << endl;
        }
        return;
    }

    for (int v = 1; v <= R[A.size()]; ++v) {
        A.push_back(v);
        dfs(A, K, R, sum + v);
        A.pop_back(); // これが結構ポイント
    }
}

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

    vector<int> A;
    dfs(A, K, R, 0);
}