【Python】4つのリストからtarget以下の合計となるユニークなインデックス組み合わせの数を効率的に求める方法

4つの整数リスト A、B、C、D とターゲット値(target)が与えられたとき、A[i] + B[j] + C[k] + D[l] ≤ target を満たすようなインデックスの組み合わせ (i, j, k, l) の総数を求める問題です。

たとえば、入力が A = [3, 2]、B = [5, 3]、C = [1]、D = [2, 3]、target = 9 の場合、出力は 3 になります。条件を満たす組み合わせとしては、[3, 3, 1, 2]、[3, 3, 1, 2]、[2, 3, 1, 3] の3通りが挙げられます。

解法のアプローチ

4つのリストすべての組み合わせを素朴に全列挙すると計算量が O(n⁴) となり、リストが大きくなると現実的な時間で処理できません。そこで、「半分ずつ処理して二分探索で結合する」という手法を使って効率化します。手順は以下のとおりです。

• まず空の一時リスト(temp_list)を用意します。
• リストAとBのすべてのペアについて、その和 A[i] + B[j] をtemp_listの末尾に追加していきます。
• temp_listを昇順にソートします。
• 続いて、リストCとDのすべてのペアについて以下を繰り返します。
  • sum_cd := C[i] + D[j](CとDのペアの和)
  • sum_ab := target − sum_cd(AとB側に許される残りの上限値)
  • ソート済みtemp_listの中から sum_ab 以下の要素の個数を数え、答え(ans)に加算します。
• 最終的なansを返します。

ここでポイントとなるのが、Python標準ライブラリ bisect モジュールの bisect_right() 関数です。ソート済みリストに対して二分探索を行い、指定した値以下の要素数を O(log n) で取得できるため、内側のカウント処理を高速化できます。

実装例

from bisect import bisect_right

class Solution:
   def solve(self, A, B, C, D, target):
      temp_list = []
      for i in range(len(A)):
         for j in range(len(B)):
            temp_list.append(A[i] + B[j])

      temp_list.sort()

      ans = 0
      for i in range(len(C)):
         for j in range(len(D)):
            sum_cd = C[i] + D[j]
            sum_ab = target - sum_cd

            ans += bisect_right(temp_list, sum_ab)

      return ans

ob = Solution()
A = [3, 2]
B = [5, 3]
C = [1]
D = [2, 3]
target = 9
print(ob.solve(A, B, C, D, target))

入力

[3, 2], [5, 3], [1], [2, 3], 9

出力

3

計算量の目安

AとBのペアの生成に O(|A|×|B|)、ソートに O(|A|×|B|・log(|A|×|B|))、さらにCとDの各ペアに対する二分探索に O(log(|A|×|B|)) かかります。全体では O((|A|×|B| + |C|×|D|)・log(|A|×|B|)) 程度に収まり、4重ループによる全列挙 O(n⁴) と比べて大幅に高速です。この「2つのグループに分けて片方をソートし、もう片方から二分探索で対応を数える」テクニックは、4Sum系の問題や「4つの配列から条件を満たす組を見つける」類題でも広く応用できるので、ぜひ覚えておきましょう。