Pythonでa×b=c×dとなる4つの異なる数の組み合わせを数える方法

問題の概要

一意な正の整数からなるリスト nums が与えられます。この中から、a × b = c × d という条件を満たし、かつ a、b、c、d がすべて互いに異なる要素であるような4つ組 (a, b, c, d) の個数を求めます。

たとえば、入力が nums = [3, 6, 4, 8] の場合、出力は 8 になります。条件を満たす4つ組は次の8通り存在するためです。

[[3,8,6,4], [3,8,4,6], [8,3,6,4], [8,3,4,6], [6,4,3,8], [4,6,3,8], [6,4,8,3], [4,6,8,3]]

解き方の考え方

すべての4つ組を総当たりで調べると計算量が O(n⁴) になり非効率です。そこで、「同じ積を持つペア」をハッシュマップ(辞書)で集計する手法を使います。手順は以下のとおりです。

• c := 空の辞書(マップ)
• n := nums の要素数
• i を 0 から n − 1 まで繰り返す:
  • j を i + 1 から n − 1 まで繰り返す:
    • x := nums[i] × nums[j]
    • c[x] := 1 +(c[x] が存在すればその値、なければ 0)
• ret := 0
• c のすべての値 x について:
  • ret := ret + x × (x − 1)
• ret × 4 を返す

なぜ最後に「× 4」するのか?

同じ積を持つペアが f 個あるとき、その中から異なる2つのペアを選ぶ方法は f × (f − 1) / 2 通りです。さらに、選んだ2つのペア (a, b) と (c, d) から4つ組を並べる方法は、「どちらのペアを先に置くか」の2通り × 各ペア内部の順序の入れ替え2通りずつ、の合計 8 通りあります。したがって全体では f × (f − 1) / 2 × 8 = f × (f − 1) × 4 通りとなり、コードではあらかじめ f × (f − 1) を合計しておいて、最後にまとめて 4 倍しています。

計算量

• 時間計算量:O(n²) ― すべてのペアの積を一度ずつ計算
• 空間計算量:O(n²) ― 最大でペアの数だけ辞書にエントリが追加される可能性

実装例(Python)

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

def solve(nums):
    c = {}
    n = len(nums)
    for i in range(n):
        for j in range(i + 1, n):
            x = nums[i] * nums[j]
            c[x] = c.get(x, 0) + 1
    ret = 0
    for x in c.values():
        ret += x * (x - 1)
    return ret * 4

nums = [3, 6, 4, 8]
print(solve(nums))

入力

[3, 6, 4, 8]

出力

8