Pythonでラベルごとの使用制限を守りながら最大値の合計を求める方法

この記事では、各アイテムが「値」と「ラベル」を持つ集合から、与えられた制約を満たしながら合計値が最大になる部分集合を見つけるアルゴリズムを、Pythonのコード例とともに解説します。

問題の概要

n個のアイテムがあるとします。i番目のアイテムは値 values[i] とラベル labels[i] を持ちます。ここから部分集合 S を選びますが、S は次の条件を満たす必要があります。

• |S| <= num_wanted（選ぶアイテムの総数は num_wanted 以下）
• どのラベル L についても、S に含まれるラベル L のアイテム数は use_limit 以下

この条件のもとで、部分集合 S の値の合計として考えられる最大値を求めます。

具体例

たとえば、values = [5,4,3,2,1]、labels = [1,1,2,2,3]、num_wanted = 3、use_limit = 1 が入力だった場合、出力は 9 になります。これは、1番目・3番目・5番目のアイテム（値 5・3・1）を選ぶことで、同じラベルを複数回使わずに合計 9 を達成できるためです。

解法のアプローチ（貪欲法）

この問題は貪欲法（グリーディ法）で効率よく解けます。ポイントは「価値の高いアイテムから順に、制約を守りながら選んでいく」ことです。以下の手順で進めます。

• (値, ラベル) のペアを格納する配列 v を作成する
• i を 0 から values の長さまでループし、v に [values[i], labels[i]] を追加する
• v を値の降順でソートする
• ans := 0、use := 空の辞書（ラベルごとの使用回数）、i := 0 で初期化する
• num_wanted が 0 より大きく、かつ i が v の長さ未満である間、以下を繰り返す
  • v[i][1]（ラベル）が use に存在しない場合：num_wanted を 1 減らし、ans に v[i][0] を加算し、use[v[i][1]] := 1 とする
  • そうでなく、use[v[i][1]] < use_limit の場合：num_wanted を 1 減らし、ans に v[i][0] を加算し、use[v[i][1]] を 1 増やす
  • i を 1 増やす
• 最後に ans を返す

値の降順にソートしておくことで、「今見ているアイテムが使えるなら必ず選ぶべき」という判断が常に正しくなります。これにより、各ラベルの使用上限とアイテム総数の両方の制約を満たしつつ、合計を最大化できます。計算量はソートが支配的となり、O(n log n) です。

実装例

それでは、実際のPythonコードを見てみましょう。

class Solution(object):
    def largestValsFromLabels(self, values, labels, num_wanted, use_limit):
        v = []
        for i in range(len(values)):
            v.append([values[i], labels[i]])
        v = sorted(v, key=lambda v: v[0], reverse=True)
        ans = 0
        use = {}
        i = 0
        while num_wanted and i < len(v):
            if v[i][1] not in use:
                num_wanted -= 1
                ans += v[i][0]
                use[v[i][1]] = 1
            elif use[v[i][1]] < use_limit:
                num_wanted -= 1
                ans += v[i][0]
                use[v[i][1]] += 1
            i += 1
        return ans

ob = Solution()
print(ob.largestValsFromLabels([5,4,3,2,1], [1,1,2,2,3], 3, 1))

入力

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

出力

9

まとめ

この問題は、アイテムを値の降順に並べ替えて先頭から順に確認し、ラベルごとの使用回数が use_limit を超えていない場合だけ選んでいく貪欲な戦略で解けます。辞書でラベルごとの使用回数を管理することで、制約チェックも O(1) で行え、全体として O(n log n) の計算量で最大の合計値を求められます。