Pythonで条件を満たす4つ組 (a, b, c, d) の個数を効率的に求めるプログラム

問題概要

数値のリスト nums が与えられます。この中から、添字について a < b < c < d を満たし、かつ nums[a] < nums[b]（前半が昇順ペア）、nums[c] > nums[d]（後半が降順ペア）という条件を満たす4つ組 (a, b, c, d) の個数を求めます。

なお、配列 nums は 1〜N の整数の順列であることが保証されています。

入力例と出力例

入力が nums = [3, 4, 7, 6, 5] の場合、答えは 5 になります。

この入力から見つかる組み合わせは、以下の5通りです。

• (3, 4, 7, 6)
• (3, 4, 6, 5)
• (3, 4, 7, 5)
• (3, 7, 6, 5)
• (4, 7, 6, 5)

解法のアプローチ

すべての4つ組を素朴に全探索すると計算量が O(N⁴) となり、非常に非効率です。そこで、次のような工夫を行います。

まず、各位置 c に対して「c より右側にある、nums[c] より小さい値の個数」（ds_smaller_than_c）を事前に計算しておきます。次に、位置 b を左から順に走査しながら、「b より左側にある、nums[b] より小さい値の個数」の累積和（as_smaller_than_b_sum）を管理します。そして、隣接する b と c = b + 1 の組み合わせごとに両者を掛け合わせて合計することで、条件を満たす4つ組の総数を効率的に求められます。

具体的な手順

1. m := 10^9 + 7（剰余演算用の定数）を定義する
2. nums の長さが 4 未満の場合は 0 を返す
3. n := nums の長さとする
4. sorted_ds := 空のリストを作成し、nums の最後の要素を挿入してソートしておく
5. ds_smaller_than_c := [0] * n で初期化する
6. c を n−2 から −1 まで 1 ずつ減らしながら繰り返す：
   • ds_smaller_than_c[c] := bisect_right(sorted_ds, nums[c] − 1)。これは sorted_ds に nums[c] − 1 を挿入してもソート順が保たれる最も右側の位置、すなわち「nums[c] より小さい要素の個数」を表す
   • sorted_ds に nums[c] をソート済みの状態で挿入する
7. quadruplet_count := 0 で初期化する
8. sorted_as := 空のリストを作成し、nums の最初の要素を挿入してソートしておく
9. as_smaller_than_b_sum := 0 で初期化する
10. b を 1 から n−2 まで繰り返す：
    • as_smaller_than_b_sum += bisect_right(sorted_as, nums[b] − 1)
    • as_smaller_than_b_sum を m で割った余りに更新する
    • sorted_as に nums[b] をソート済みの状態で挿入する
    • quadruplet_count += as_smaller_than_b_sum * ds_smaller_than_c[b + 1]
    • quadruplet_count を m で割った余りに更新する
11. quadruplet_count を返す

実装例

それでは、実際の実装を見てみましょう。Python の bisect モジュールを使うことで、二分探索による挿入位置の検索と、ソート順を保ったままの挿入を簡潔に記述できます。

import bisect

MOD = 10 ** 9 + 7

class Solution:
    def solve(self, nums):
        if len(nums) < 4:
            return 0
        n = len(nums)

        # 各位置cより右側にある、nums[c]より小さい値の個数を事前計算
        sorted_ds = [nums[-1]]
        ds_smaller_than_c = [0] * n
        for c in range(n - 2, -1, -1):
            ds_smaller_than_c[c] = bisect.bisect_right(sorted_ds, nums[c] - 1)
            bisect.insort(sorted_ds, nums[c])

        # 各位置bより左側にある、nums[b]より小さい値の個数の累積和を利用
        quadruplet_count = 0
        sorted_as = [nums[0]]
        as_smaller_than_b_sum = 0
        for b in range(1, n - 2):
            as_smaller_than_b_sum += bisect.bisect_right(sorted_as, nums[b] - 1)
            as_smaller_than_b_sum %= MOD
            bisect.insort(sorted_as, nums[b])
            quadruplet_count += as_smaller_than_b_sum * ds_smaller_than_c[b + 1]
            quadruplet_count %= MOD

        return quadruplet_count

ob = Solution()
print(ob.solve([3, 4, 7, 6, 5]))

入力

[3, 4, 7, 6, 5]

出力

5

計算量について

bisect による挿入位置の検索は O(log N)、リストへの挿入は O(N) かかるため、この実装全体の計算量は O(N²) となります。素朴な全探索の O(N⁴) と比べると大幅な高速化であり、N が数千程度の規模でも十分に実用的です。