Pythonで異なる種類のアイテムから作れるサイズKのグループの最大数を求める方法

問題の概要

数値のリスト counts が与えられ、counts[i] は種類 i のアイテムの個数を表しているとします。さらに別の値 k も与えられます。このとき、各グループが必ず異なる種類のアイテムのみで構成されるという条件のもとで、サイズ k のグループを最大いくつ作れるかを求めるのが本記事の目的です。

具体例

たとえば、入力が counts = [2, 3, 5, 3]、k = 2 の場合、出力は 6 になります。

4種類のアイテムをそれぞれ a、b、c、d と表すと、すべての要素が異なる種類であるようなサイズ2のグループとして、次の6つを作ることができます。

[(c, a), (b, a), (c, b), (c, b), (d, a), (d, a)]

解き方のアプローチ

この問題は「二分探索」と「貪欲法（グリーディ法）」を組み合わせることで効率的に解けます。作成できるグループ数 g を仮定したとき、「g 個のグループが実際に作れるか」を判定する関数 possible() を用意し、その判定結果をもとにグループ数を二分探索します。

possible() 関数のロジック

• 必要なアイテム総数 required を「groups × k」として初期化する。
• 各種類 i について、その種類から使えるアイテム数 temp を「min(counts[i], groups, required)」とする。同じグループ内に同種のアイテムを複数入れられないため、各種類からは最大 groups 個までしか使用できない点がポイントです。
• required から temp を差し引き、required が 0 になれば True を返す。
• ループが終了しても required が 0 にならなければ False を返す。

solve() 関数のロジック

• 探索範囲の下限 l を 0、上限 r を counts の合計値とする。
• l ≤ r の間、中央値 m を計算し、possible(counts, m, k) が真なら答えの候補として記録し、探索範囲を右側へ広げる。
• 偽であれば探索範囲を左側へ狭める。
• ループ終了後の res を結果として返す。

実装例

以下の実装を見ると、処理の流れがより理解しやすくなります。

def possible(counts, groups, k):
    required = groups * k
    for i in range(len(counts)):
        temp = min(counts[i], groups, required)
        required -= temp
        if required == 0:
            return True
    return False

def solve(counts, k):
    res = 0
    l = 0
    r = sum(counts)
    while l <= r:
        m = l + (r - l) // 2
        if possible(counts, m, k):
            l = m + 1
            res = max(res, m)
        else:
            r = m - 1
    return res

counts = [2, 3, 5, 3]
k = 2
print(solve(counts, k))

入力

[2, 3, 5, 3], 2

出力

6

計算量について

possible() の判定は O(n)（n はアイテムの種類数）、二分探索は O(log S)（S はアイテムの総数）回実行されるため、全体の計算量は O(n log S) となります。考えられるグループ数を1つずつ全探索する方法に比べて格段に効率的であり、大規模な入力にも対応できる実用的な手法です。