PythonでK個の最大合計ペアを見つけるプログラム

問題の概要

2つの数値リスト nums0 と nums1、および整数 k が与えられたとします。目標は、nums0 の要素1つと nums1 の要素1つから構成されるペアの中から、合計値が大きい方から k 個のペアを見つけ出し、それらの合計を返すことです。

例えば、入力が nums1 = [8, 6, 12]、nums2 = [4, 6, 8]、k = 2 の場合、出力は 38 になります。これは、最大のペアが [12, 8](合計20)と [12, 6](合計18)であり、20 + 18 = 38 となるためです。

解法のアプローチ

この問題は、ヒープ(優先度付きキュー)を使うことで効率的に解けます。全てのペアの組み合わせを生成すると計算量が膨大になりますが、大きな合計を持つペアから順番に取り出すことで、必要な分だけ処理できます。具体的な手順は以下の通りです。

• k > len(nums0) * len(nums1) の場合(作成可能なペアの総数より多く要求された場合)、0を返します。
• 最小ヒープ pq を新しく作成します。
• 答えを格納する変数 ans を0で初期化します。
• nums0 と nums1 をそれぞれ昇順にソートします。
• i, j を、それぞれ両リストの末尾のインデックス(サイズ − 1)に設定します。
• 訪問済みのインデックス組み合わせを記録するためのセット visited を作成します。
• ヒープ pq に (−(nums0[i] + nums1[j]), i, j) をプッシュします。
• 0から k 回までループを実行し、以下の処理を行います。
  • ヒープ pq から sum, i, j をポップします。
  • x := nums0[i − 1] + nums1[j] を計算します。
  • (i − 1, j) が visited に存在しない場合は、visited に追加し、ヒープに (−x, i − 1, j) をプッシュします。
  • y := nums0[i] + nums1[j − 1] を計算します。
  • (i, j − 1) が visited に存在しない場合は、visited に追加し、ヒープに (−y, i, j − 1) をプッシュします。
  • ans := ans + (−sum) として答えを累積します。
• 最後に ans を返します。

アルゴリズムのポイント

この手法の鍵となるのは、次の2点です。

• 負の値による最大ヒープのシミュレーション: Pythonの heapq モジュールは最小ヒープしか提供していません。そこで、合計値にマイナスを付けて格納することで、最小ヒープを最大ヒープとして振る舞わせています。
• 隣接ペアへの段階的な展開: ソート済みリストの末尾同士のペアが最大の合計となります。ヒープからペアを取り出すたびに、「片方のリストのインデックスを1つ左にずらした隣接ペア」だけを候補として追加します。visited セットで重複を防ぐことで、無駄な計算を避けています。

これにより、全ペアを生成することなく、上位 k 個のペアだけを効率よく求められます。

Pythonでの実装例

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

from heapq import heappush, heappop
class Solution:
    def solve(self, nums0, nums1, k):
        if k > len(nums0) * len(nums1):
            return 0
        pq = []
        ans = 0
        nums0.sort(), nums1.sort()
        i, j = len(nums0) - 1, len(nums1) - 1
        visited = set()
        heappush(pq, (-(nums0[i] + nums1[j]), i, j))
        for _ in range(k):
            sum, i, j = heappop(pq)
            x = nums0[i - 1] + nums1[j]
            if not (i - 1, j) in visited:
                visited.add((i - 1, j))
                heappush(pq, (-x, i - 1, j))
            y = nums0[i] + nums1[j - 1]
            if not (i, j - 1) in visited:
                visited.add((i, j - 1))
                heappush(pq, (-y, i, j - 1))
            ans += -sum
        return ans
ob = Solution()
print(ob.solve([8, 6, 12], [4, 6, 8], 2))

入力

[8, 6, 12],[4, 6, 8],2

出力

38