Pythonですべての学生グループを収容できるバスのサイズを求めるプログラム

問題の概要

n個の学生グループが、大学バスで自宅に帰るため待機しています。各グループにはm人の学生がいます。学生たちは仲間と離ればなれになることを嫌うため、グループの全員が乗車できる場合にのみバスに乗ります。さらに、グループは必ず順番を守って乗車し、前のグループが乗り終わる(または目的地に到着する)までは自分の順番が回ってきません。

そこで、グループの数と各グループの人数が与えられたとき、「すべてのグループを輸送でき、かつ大学を出発するたびにバス内に空席がひとつもない」ようなバスのサイズをすべて求めます。

入力例と出力例

たとえば、各グループの人数が gr_no = [3, 4, 2, 2, 1, 4, 3, 5] の場合、出力は [12, 24] になります。

• バスのサイズが12の場合: 1回目の運行でグループ1〜5(3+4+2+2+1=12人)を乗せ、2回目の運行で残りのグループ(4+3+5=12人)を運べます。
• バスのサイズが24の場合: 全グループ(合計24人)を1回の運行でまとめて運べます。

解法のアプローチ

この問題は「総人数の約数」に注目すると効率的に解けます。毎回の運行でバスが満席になるということは、運行の区切りとなる累積人数がバスのサイズの整数倍になっている必要があるためです。手順は以下の通りです。

1. 約数を列挙する関数 factor_ret() を定義します。 引数 n の平方根までの整数 i を順に確認し、n が i で割り切れる場合は i と n//i を約数として収集します。最後に重複を除去し、昇順にソートして返します。
2. 累積和を計算します。 total[i] に「先頭から i 番目のグループまでの合計人数」を格納します。この値は、その時点までに乗車した人数を表します。
3. 各約数について妥当性を検証します。 総人数の約数 size のそれぞれについて、累積和の中で size で割り切れる値だけを取り出し、それらが size×1, size×2, … と連続した倍数列になっていれば、size は条件を満たすバスのサイズです。
4. 条件を満たすすべてのサイズを結果リストに追加して返します。

実装例(Python)

以下のコードで実際の実装を確認できます。

from functools import reduce

def solve(gr_no):
    # 累積和を計算
    total = [gr_no[0]]
    for i in range(1, len(gr_no)):
        total.append(total[i - 1] + gr_no[i])
    b_sizes = []
    # 総人数の各約数について検証
    for size in factor_ret(sum(gr_no)):
        # size で割り切れる累積和のみ抽出
        temp_list = list(filter(lambda x: x % size == 0, total))
        index = 1
        indicator = True
        for point in temp_list:
            if point != size * index:
                indicator = False
                break
            index += 1
        if indicator:
            b_sizes.append(size)
    return b_sizes

def factor_ret(n):
    # n の約数をすべて列挙(重複なし・昇順)
    return sorted(set(reduce(list.__add__,
        ([i, n // i] for i in range(1, int(n**0.5) + 1) if n % i == 0))))

print(solve([3, 4, 2, 2, 1, 4, 3, 5]))

実行結果

入力:

[3, 4, 2, 2, 1, 4, 3, 5]

出力:

[12, 24]

アルゴリズムのポイント

バスが毎回満席で出発するためには、運行の境界となる累積和が必ずバスサイズの倍数である必要があります。逆に言えば、累積和のうちサイズの倍数となっている地点が size×1, size×2, … という連続的な倍数列を形成していれば、そのサイズでグループを区切って満席の運行を組むことができます。約数列挙は O(√N)、各約数の検証は累積和の走査だけで済むため、全体として非常に効率的な解法となっています。