【Python】配列要素とインデックスが一致する最小のインデックスを見つけるプログラム

問題の概要

すべての要素が重複なく、昇順にソートされたリスト nums が与えられます。この中から、nums[i] = i(要素の値とインデックスが一致する)を満たす最小のインデックス i を見つける必要があります。条件を満たすインデックスが存在しない場合は -1 を返してください。さらに、この問題は O(log n) の時間計算量で解くことが求められています。

たとえば、入力が nums = [-4, -1, 2, 3, 8] の場合、出力は 2 となります。nums[2] = 2 と nums[3] = 3 の両方が条件を満たしていますが、より小さい方の「2」が答えになるためです。

解法の考え方:二分探索

O(log n) という計算量の制約を満たすには、先頭から順に調べる線形探索ではなく、二分探索(バイナリサーチ)を活用します。配列が昇順にソートされ、かつ要素がすべて一意であることを利用すると、次のような性質が成り立ちます。

• nums[mid] >= mid の場合:それより右側のインデックスでは必ず値がインデックスを上回るため、答えは左半分にしか存在しません。探索範囲を左側に狭めます。
• nums[mid] < mid の場合:それより左側に条件を満たすインデックスは存在しないため、探索範囲を右半分に狭めます。

この性質により、各ステップで探索範囲を半分に絞り込みながら、効率的に答えを求めることができます。

アルゴリズムの手順

• 変数 ret を -1 で初期化し、探索範囲の左端 lhs を 0、右端 rhs を「リストの長さ − 1」とします。
• lhs <= rhs の間、以下を繰り返します。
  • mid を (lhs + rhs) // 2(切り捨て除算)として計算します。
  • nums[mid] == mid であれば、ret = mid として答えを記録します。
  • nums[mid] >= mid であれば、rhs = mid - 1 として左側を探索します。
  • そうでなければ、lhs = mid + 1 として右側を探索します。
• ループ終了後、ret を返します。

Pythonでの実装例

以下が実際の実装コードです。

def solve(nums):
   ret = -1
   lhs = 0
   rhs = len(nums) - 1
   while lhs <= rhs:
      mid = (lhs + rhs) // 2
      if nums[mid] == mid:
         ret = mid
      if nums[mid] >= mid:
         rhs = mid - 1
      else:
         lhs = mid + 1
   return ret

nums = [-4, -1, 2, 3, 8]
print(solve(nums))

入力

[-4, -1, 2, 3, 8]

出力

2

計算量のまとめ

• 時間計算量: O(log n) — 各反復で探索範囲が半分になるため。
• 空間計算量: O(1) — 追加のデータ構造を必要としないため。