Pythonでリストの各要素より右側にある小さい要素の数を返すプログラム

問題概要

数値のリスト nums が与えられたとします。このとき、元のリストの各要素について、その要素より右側に存在する小さい要素の個数を求め、それらを並べた新しいリストを作成する問題を考えてみましょう。

例えば、入力が nums = [4, 5, 9, 7, 2] の場合、出力は [1, 1, 2, 1, 0] になります。これは次のような理由によるものです。

• 4 の右側にある小さい要素は 1 個(2)
• 5 の右側にある小さい要素は 1 個(2)
• 9 の右側にある小さい要素は 2 個(7 と 2)
• 7 の右側にある小さい要素は 1 個(2)
• 2 の右側には小さい要素が存在しないため 0 個

解法のアプローチ

この問題は、リストを右から左へ走査しながら、これまでに登場した要素をソート済みの状態で管理することで効率的に解けます。具体的な手順は以下の通りです。

1. 結果を格納するリスト res と、ソート済みリスト inc を新しく作成します。
2. nums が空になるまで、次の処理を繰り返します。
   • nums の末尾から要素 num を取り出します。
   • inc に num を挿入すべき位置(左端のインデックス)を調べ、その値を res の末尾に追加します。この挿入位置こそが「右側にある小さい要素の数」に相当します。
   • num を inc に挿入し、ソート済みの状態を維持します。
3. 最後に res を逆順にして返します(末尾から処理を進めたため)。

ここでポイントとなるのが、Python標準ライブラリの bisect モジュールです。bisect_left() は二分探索により挿入位置を O(log n) で求められ、insort() はソート順を崩さずに要素を挿入できるため、全体の計算量を O(n²) ではなく O(n log n) 程度に抑えられます。

実装例

それでは、実際のコードを見てみましょう。

import bisect
class Solution:
    def solve(self, nums):
        res, inc = [], []
        while nums:
            num = nums.pop()
            res.append(bisect.bisect_left(inc, num))
            bisect.insort(inc, num)
        return res[::-1]
ob = Solution()
nums = [4, 5, 9, 7, 2]
print(ob.solve(nums))

コードの解説

• nums.pop():リストの末尾から要素を取り出し、右側から順に処理していきます。
• bisect.bisect_left(inc, num):ソート済みリスト inc の中で num より小さい要素の個数(= 挿入位置)を取得します。
• bisect.insort(inc, num):num を適切な位置に挿入し、inc を常にソート済みに保ちます。
• res[::-1]:右から処理した結果を反転させ、元のリストの順序に合わせて返します。

入力

[4, 5, 9, 7, 2]

出力

[1, 1, 2, 1, 0]

まとめ

本記事では、リストの各要素について右側にある小さい要素の数を求めるPythonプログラムを紹介しました。bisect モジュールの二分探索を活用することで、単純な全比較(O(n²))よりも高速に処理でき、大規模なデータにも対応できます。同様の「右側・左側の大小カウント」系の問題では、今回のようなソート済みリスト管理のテクニックが非常に役立ちますので、ぜひ覚えておきましょう。