【Python】リストからn番目に小さい要素を線形時間で選択するプログラムの書き方

はじめに

線形時間計算量 O(n) でリストから n 番目に小さい要素を選択するには、大きく分けて2つの処理が必要になります。ひとつはピボット(基準値)を決めてリストを分割する「パーティション」処理、もうひとつはその結果をもとに目的の要素を再帰的に絞り込んでいく「選択」処理です。

リストの分割は、ユーザーが指定した「i」の値に基づいて行われます。この値とピボット位置との大小関係を比較することで、探索範囲を左半分か右半分のどちらかに限定でき、不要な計算を省けます。この手法は「クイックセレクト」と呼ばれるアルゴリズムで、平均計算量は O(n) となるのが特徴です。

以下に実際の実装例を示します。

サンプルコード

def select_smallest(my_list, beg, end, i):
    if end - beg <= 1:
        return my_list[beg]
    pivot_val = start_partition(my_list, beg, end)

    k = pivot_val - beg + 1

    if i < k:
        return select_smallest(my_list, beg, pivot_val, i)
    elif i > k:
        return select_smallest(my_list, pivot_val + 1, end, i - k)

    return my_list[pivot_val]

def start_partition(my_list, beg, end):
    pivot_val = my_list[beg]
    i = beg + 1
    j = end - 1

    while True:
        while (i <= j and my_list[i] <= pivot_val):
            i = i + 1
        while (i <= j and my_list[j] >= pivot_val):
            j = j - 1

        if i <= j:
            my_list[i], my_list[j] = my_list[j], my_list[i]
        else:
            my_list[beg], my_list[j] = my_list[j], my_list[beg]
            return j

my_list = input('Enter the list of numbers.. ')
my_list = my_list.split()
my_list = [int(x) for x in my_list]
i = int(input('Enter the value for i.. '))

ith_smallest = select_smallest(my_list, 0, len(my_list), i)
print('The result is {}.'.format(ith_smallest))

実行結果

Enter the list of numbers.. 43 12 67 89 99 0
Enter the value for i.. 3
The result is 43.

コードの解説

• まず「select_smallest」という関数を定義します。引数として、対象のリスト、開始インデックス(beg)、終了インデックス(end)、そして「i」の値を受け取ります。

• 次に「start_partition」という関数を定義します。この関数は、先頭の要素をピボットとして扱い、「i」の値に応じてリストを2つの部分に分割します。

• 「start_partition」は「select_smallest」の中から呼び出され、ピボットの最終的な位置を返します。

• 「select_smallest」は条件に応じて自分自身を再度呼び出します。これが再帰処理であり、探索範囲が目的の要素1つに絞り込まれるまで繰り返されます。

• 数値のリストは、input() 関数によってユーザーからの入力として受け取ります。

• 受け取った文字列は、split() メソッドによりデフォルトの区切り文字(空白)を基準に分割されます。

• リスト内包表記を使って各要素を整数(int)に変換しながら反復処理します。

• 「i」の値も同様にユーザーから入力として受け取ります。

• この「i」の値とピボット位置の関係に基づき、リストは2つの部分に分割されます。

• 目的の要素が存在する側の部分リストに対してのみ「select_smallest」が再帰的に呼び出されるため、全体の計算量を抑えられます。

• 最終的に求まった n 番目に小さい値が、コンソールに出力されます。

計算量について

このアルゴリズムの平均時間計算量は O(n) です。各再帰ステップで探索対象がおおよそ半分になるためです。ただし、ピボットの選び方が運悪く偏った場合、最悪計算量は O(n²) になる点には注意が必要です。実用上は、ピボットをランダムに選ぶことで最悪ケースの発生確率を大幅に下げられます。

まとめ

リスト全体をソートしてから n 番目の要素を取り出す方法(O(n log n))と比べ、クイックセレクトを用いれば平均 O(n) で n 番目に小さい要素を取得できます。「パーティションによる分割」と「再帰的な絞り込み」という2つの処理を組み合わせることが、このアルゴリズムの核心です。順序統計量を効率的に求めたい場面でぜひ活用してみてください。