C++で円状に並んだ箱からすべての石を取り除けるか判定する方法

問題の概要

N個の要素を持つ配列Aがあるとしましょう。N個の箱が円状に並べられており、i番目の箱にはA[i]個の石が入っています。以下の操作を繰り返し行うことで、すべての箱から石を取り除くことができるかどうかを判定する必要があります。

操作の内容： 箱を1つ選択します(これをi番目の箱と呼びます)。jを1からNまで変化させながら、各ステップで(i+j)番目の箱からちょうどj個の石を取り除きます。なお、(N+k)番目の箱はk番目の箱として扱われます。選ばれた箱に十分な数の石が残っていない場合、この操作は実行できません。

例えば、入力が A = [4, 5, 1, 2, 3] の場合、出力は True となります。2番目の箱から操作を開始すると、すべての石を取り除くことができるためです。

解法のアプローチ

この問題を解くために、以下の手順に従います。

n := 配列Aのサイズ
サイズ(n+1)の配列aを定義
サイズ(n+1)の配列bを定義
sum := 0, p := n * (n + 1) / 2
iを1からnまで繰り返す:
   a[i] := A[i - 1]
   sum := sum + a[i]
sum mod p が 0 でない場合:
   return false
k := sum / p
iを1からnまで繰り返す:
   b[i] := a[i] - a[(i mod n) + 1]
sum := 0
iを1からnまで繰り返す:
   a[i] := b[i]
   sum := sum + a[i]
sum が 0 でない場合:
   return false
iを1からnまで繰り返す:
   (a[i] + k) mod n が 0 でない、または a[i] + k < 0 の場合:
      return false
return true

アルゴリズムのポイント

1回の操作で取り除ける石の総数は、1+2+…+N = N(N+1)/2 個です。そのため、まず全体の石の総数が N(N+1)/2 で割り切れることを確認し、割り切れる場合はその商 k が「操作を行う回数」に相当します。

次に、隣接する箱同士の石の個数の差 a[i] − a[i+1] を計算します。これらの差の総和は必ず0になる必要があります。さらに、各差について (a[i] + k) が n で割り切れ、かつ負でないことも確認します。これらの条件をすべて満たす場合のみ、すべての石を取り除くことが可能です。

C++での実装例

理解を深めるために、以下の実装例を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
bool solve(vector<int> A) {
    int n = A.size();
    vector<int> a(n + 1);
    vector<int> b(n + 1);
    int sum = 0, p = n * (n + 1) / 2;
    for (int i = 1; i <= n; i++) {
        a[i] = A[i - 1];
        sum += a[i];
    }
    if (sum % p != 0) {
        return false;
    }
    int k = sum / p;
    for (int i = 1; i <= n; i++) {
        b[i] = a[i] - a[i % n + 1];
    }
    sum = 0;
    for (int i = 1; i <= n; i++) {
        a[i] = b[i];
        sum += a[i];
    }
    if (sum != 0) {
        return false;
    }
    for (int i = 1; i <= n; i++) {
        if ((a[i] + k) % n != 0 || a[i] + k < 0) {
            return false;
        }
    }
    return true;
}
int main(){
    vector<int> A = { 4, 5, 1, 2, 3 };
    cout << solve(A) << endl;
}

入力

{ 4, 5, 1, 2, 3 }

出力

1