Pythonで2つのリストの要素の積からk番目に大きい値を求めるプログラム

整数を含む2つのリスト p と q が与えられたとします。この問題では、両リストの要素をすべて掛け合わせた積の中から、k番目に大きい値(インデックスは0から開始)を求めます。

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

考えられる積は次の4つです。
2 × 6 = 12
2 × 8 = 16
5 × 6 = 30
5 × 8 = 40
これらを降順に並べると [40, 30, 16, 12] となるため、インデックス2の要素、すなわち求める答えは 16 です。

解法のアプローチ

すべての組み合わせの積を計算してソートする方法もありますが、要素数が多い場合は非効率になります。そこで、最小ヒープ(min-heap)を活用し、「上位の積だけを常に保持する」戦略でこの問題を解きます。

具体的には、サイズ k+1 のヒープを維持しながら積を順に挿入していき、最終的にヒープの先頭(最小値)が求める「インデックスkの要素」になります。また、ヒープが満杯の状態で現在の積がヒープの最小値以下であれば、それ以降の積は必ずさらに小さくなるため、探索を途中で打ち切る枝刈りを行うことで処理を高速化できます。

アルゴリズムの手順

1. リスト p をソートします。
2. リスト q をソートします。
3. k := k + 1 とします。
4. heap := リストで表現した新しい空のヒープを用意します。
5. q の各要素 elem について、次の処理を行います。
   • elem ≥ 0 の場合: i を p のサイズ - 1 から -1 まで1ずつ減らしながら繰り返します。
     • cd := elem * p[i] を計算します。
     • heap が空でなく、heap のサイズが k と等しく、cd ≤ heap[0] である場合はループを抜けます。
     • cd を heap に挿入します。
     • heap の長さが k を超えたら、heap から最小の要素を削除します。
   • elem < 0 の場合: i を 0 から p のサイズまで1ずつ増やしながら、上記と同じ処理を繰り返します。
6. 最後に heap[0] を返します。

ここで重要なのは走査の方向です。elem が正の数なら、p の大きい要素との積ほど大きくなるため降順に走査し、elem が負の数なら、p の小さい要素(絶対値が大きい負数)との積ほど大きくなるため昇順に走査します。これにより、大きい積から順に評価できるので、早期打ち切りの判定が効果的に機能します。

実装例

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

from heapq import heappush, heappop

def solve(p, q, k):
   p = sorted(p)
   q = sorted(q)
   k += 1
   heap = []
   for elem in q:
      if elem >= 0:
         for i in range((len(p) - 1), -1, -1):
            cd = elem * p[i]
            if heap and len(heap) == k and cd <= heap[0]:
               break
            heappush(heap, cd)
            if len(heap) > k:
               heappop(heap)
      else:
         for i in range(len(p)):
            cd = elem * p[i]
            if heap and len(heap) == k and cd <= heap[0]:
               break
            heappush(heap, cd)
            if len(heap) > k:
               heappop(heap)
   return heap[0]

print(solve([2, 5], [6, 8], 2))

入力

[2, 5], [6, 8], 2

出力

16

まとめ

全組み合わせの積(n × m 個)をすべて生成してソートする方法(O(nm log nm))と比べ、このアルゴリズムはヒープによる枝刈りによって、多くのケースで大幅に少ない計算量で答えを得られます。ソートと走査方向を工夫して大きい積から順に評価できるようにしている点が、この手法の大きなポイントです。