Pythonでリストからn番目に大きい要素を線形時間で選択するプログラム

リストから線形時間計算量 O(n) で n 番目に大きい要素を選択するには、主に2つの処理が必要になります。

• ピボットを基準にリストを2つの部分に分割する「パーティション」処理
• 分割された部分の中から目的の要素を再帰的に絞り込んでいく「選択」処理

この手法は、クイックソートの考え方を応用した Quickselect（クイックセレクト）アルゴリズム として知られています。ユーザーが指定した「i」の値に基づいてリストを分割し、そのたびに探索範囲を半分以下に狭めていくことで、全体をソートせずに効率よく目的の要素を見つけられます。

サンプルコード

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

   k = end - pivot_val

   if i < k:
      return select_largest(my_list, pivot_val + 1, end, i)
   elif i > k:
      return select_largest(my_list, beg, pivot_val, 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_largest = select_largest(my_list, 0, len(my_list), i)
print('The result is {}.'.format(ith_largest))

実行結果

Enter the list of numbers.. 34 67 12 0 999
Enter the value for i.. 1
The result is 999.

コードの解説

• select_largest メソッドが定義され、リスト・開始インデックス・終了インデックス・「i」の値を引数として受け取ります。

• start_partition メソッドが定義され、ピボット値を基準にリストを2つの部分に分割します。

• パーティション処理は select_largest の内部から呼び出されます。

• select_largest は自分自身を再度呼び出しています。これが再帰処理による実装です。

• 数値のリストはユーザーからの入力として受け取ります。

• split() によって入力文字列が空白区切りで分割されます。

• リスト内包表記により、各要素が整数へと変換されます。

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

• この「i」の値に基づき、リストが2つの部分に振り分けられます。

• 目的の要素が存在する側の部分リストに対してのみ select_largest を再帰的に適用し、不要な側は捨てることで計算量を抑えています。

• 最終的な結果がコンソールに出力されます。

計算量について

平均的なケースでは、各再帰ステップで探索対象が半分程度に減っていくため、O(n) + O(n/2) + O(n/4) + … = O(2n) ≒ O(n) となり、線形時間での選択が可能です。ただし、ピボットの選び方が不運な場合（最悪ケース）には O(n²) まで悪化する可能性がある点には注意が必要です。