Pythonでリストの要素を昇順に削除した際のインデックスを取得するプログラム

問題の概要

重複のない値を持つリストが与えられ、各数値を小さい方から順番(昇順)に削除していくことを考えます。このとき、各数値が削除された時点でのインデックスを、削除の順序どおりに求めるのがこの問題の目的です。

例えば、入力が nums = [4, 6, 2, 5, 3, 1] の場合、出力は [5, 2, 3, 0, 1, 0] になります。処理の流れは以下のとおりです。

• 最初に 1 を削除 → 配列は [4, 6, 2, 5, 3](インデックス 5)
• 次に 2 を削除 → 配列は [4, 6, 5, 3](インデックス 2)
• 次に 3 を削除 → 配列は [4, 6, 5](インデックス 3)
• 次に 4 を削除 → 配列は [6, 5](インデックス 0)
• 次に 5 を削除 → 配列は [6](インデックス 1)
• 最後に 6 を削除 → 配列は空になる(インデックス 0)

解法のポイント

実は、ある値が削除されるときのインデックスは、元のリスト内でその値より大きく、かつ前方に位置する要素の個数と一致します。上の例で値 3 に注目すると、前方にあるより大きい要素は 4・6・5 の 3 つなので、答えは 3 となります。

この「前方にある大きい要素の個数」は、マージソートを応用することで効率的に数えられます。マージ処理の際に右側の要素が先に配置されるとき、まだ配置されていない左側の要素はすべてその値より大きいため、その個数を記録していきます。

アルゴリズムの手順

• 関数 my_sort() を定義します。引数としてインデックスのリスト inds を受け取ります。
• inds のサイズが 1 以下の場合は、そのまま返します。
• 新しいリスト sorted_inds を用意します。
• mid := len(inds) // 2 として中央位置を求めます。
• left := my_sort(inds[:mid])、right := my_sort(inds[mid:]) として、左右を再帰的にソートします。
• i = 0、j = 0 と初期化します。
• i < len(left) かつ j < len(right) の間、以下を繰り返します。
  • nums[left[i]] < nums[right[j]] の場合:sorted_inds の末尾に left[i] を追加し、i を 1 増やします。
  • それ以外の場合:sorted_inds の末尾に right[j] を追加し、larger[right[j]] に len(left) - i を加算して、j を 1 増やします。
• 残りの left[i:] と right[j:] を sorted_inds に連結します。
• sorted_inds を返します。

メイン処理の流れ

• larger:nums と同じサイズの、0 で初期化されたリストを用意します。
• my_sort(range(len(nums))) を呼び出してカウントを行います。
• num_larger_pairs:nums と larger の各要素をペアにしてソートします。
• num_larger_pairs の各要素 e から e[1] を取り出したリストを返します。

実装例(Python)

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

class Solution:
    def solve(self, nums):
        return solve(nums)

def solve(nums):
    def my_sort(inds):
        if len(inds) <= 1:
            return inds
        sorted_inds = []
        mid = len(inds) // 2
        left, right = my_sort(inds[:mid]), my_sort(inds[mid:])
        i = j = 0
        while i < len(left) and j < len(right):
            if nums[left[i]] < nums[right[j]]:
                sorted_inds.append(left[i])
                i += 1
            else:
                sorted_inds.append(right[j])
                larger[right[j]] += len(left) - i
                j += 1
        sorted_inds.extend(left[i:])
        sorted_inds.extend(right[j:])
        return sorted_inds
    larger = [0] * len(nums)
    my_sort(range(len(nums)))
    num_larger_pairs = sorted(zip(nums, larger))
    return [e[1] for e in num_larger_pairs]

ob = Solution()
nums = [4, 6, 2, 5, 3, 1]
print(ob.solve(nums))

入力

[4, 6, 2, 5, 3, 1]

出力

[5, 2, 3, 0, 1, 0]

計算量について

このアルゴリズムはマージソートと同じ再帰構造を持つため、時間計算量は O(n log n)、空間計算量は O(n) となります。毎回最小値を探して削除をシミュレーションする素朴な手法(O(n²))と比べ、要素数が多い入力でも高速に動作するのが大きな利点です。