B - Tetrahedral Number / AtCoder Beginner Contest 335（Sponsored by Mynavi）

問題

整数 が与えられる。

非負整数の組 であって を満たすものを辞書順で小さい方から全て出力せよ。

入力

1行に、整数 が与えられる。

条件

• 実行時間制限: 2s
• メモリ制限: 1024MB
• 

出力

条件を満たす の組を、1行に1組ずつ辞書順で小さい方から順にすべて出力すること。

解法

条件を満たすように の値が既に決定されているとすると、残る の値の候補として考えられるのは です。

このとき、 の値が小さい方から順に出力することで、これらを辞書順で小さい方から出力するということになります。

同様にして、 の値が既に決定されているとすると、辞書順で小さい方から出力するには、

• であるものを、辞書順で小さい方から出力する
• であるものを、辞書順で小さい方から出力する
• 
• であるものを、辞書順で小さい方から出力する

という手順によって、辞書順で小さい方からの出力が達成されます。

もちろん、 についても同様のことが言えるので、全体の手順としては

• と設定する
  • と設定する
    • と設定し、 を出力する
    • と設定し、 を出力する
    • 
    • と設定し、 を出力する
  • と設定する
    • と設定し、 を出力する
    • と設定し、 を出力する
    • 
    • と設定し、 を出力する
  • 
  • と設定する
    • と設定し、 を出力する
• と設定する
  • と設定する
    • と設定し、 を出力する
    • と設定し、 を出力する
    • 
    • と設定し、 を出力する
  • と設定する
    • と設定し、 を出力する
    • と設定し、 を出力する
    • 
    • と設定し、 を出力する
  • 
  • と設定する
    • と設定し、 を出力する
• 
• と設定する
  • と設定する
    • と設定し、 を出力する

というものになります。

これは、3重のforループを書くことにより実装ができます。 また、全体の計算量としても 程度であり、 であるので、実行時間制限にも十分間に合います。

ソースコード

main() 関数の中に、答えを出力する部分を直接実装しました。

int main() {
  // 値の入力
  int n;
  cin >> n;

  // 3重ループを用いて、条件を満たす組み合わせを出力していく
  for (int i = 0; i <= n; i++) {
    for (int j = 0; i + j <= n; j++) { // 条件を i + j <= n とすることで、必ず条件が満たされるようにする
      for (int k = 0; i + j + k <= n; k++) { // 条件を i + j + k <= n とすることで、必ず条件が満たされるようにする
        cout << i << " " << j << " " << k << endl; // 組み合わせの出力
      }
    }
  }

  return 0;
}

• 提出したコード: https://atcoder.jp/contests/abc335/submissions/49364625
• GitHub

再帰を用いるソースコード

今回の問題では3重ループで済むので特に問題にはなりませんが、整数の個数が多い場合などについては再帰を用いることによって、forループの多用を防ぐことができます。

引数として「現在何個目の値について見ているのか」を表す整数（コード上では turn ）と、「値として使用できるものの最大値」を表す整数（コード上では rest ）の2つを持つ関数 solve() を以下のように実装しました。

int n;
int a[3]; // a は (x, y, z) の値が入る配列

void solve(int turn, int rest) {
  if (turn == 3) {
    // turn が 3 となったら a の中身の値をすべて出力して終了する
    cout << a[0] << " " << a[1] << " " << a[2] << endl;
    return;
  }

  // 0 から rest までの整数を順番にチェックする
  for (int i = 0; i <= rest; i++) {
    a[turn] = i; // a[turn] に i を入れる
    solve(turn + 1, rest - i); // turn を 1 増やし、 rest から turn を引いた状態で、再帰的に solve() 関数に入れる
  }
  return;
}

この solve() 関数では、最終的に値を出力する部分までも含めており、また各 turn において、値の小さい順に配列 a に保存していることから、これで再帰的に辞書順の表示が行えると言えます。

• 提出したコード: https://atcoder.jp/contests/abc335/submissions/49348876
• GitHub

感想

実は、最初に思いついたのが再帰的に書く方法だったので、AtCoderのB問題の難易度では全くないなと思っていました。

3重ループにする方法については、この記事を書きながら思いついたものです。 単純に考えることが重要ですね。