LeetCode / Search Insert Position - オットセイの経営日誌

https://leetcode.com/problems/search-insert-position/

Given a sorted array and a target value, return the index if the target is found. If not, return the index where it would be if it were inserted in order.

You may assume no duplicates in the array.

Example 1:
Input: [1,3,5,6], 5
Output: 2

Example 2:
Input: [1,3,5,6], 2
Output: 1

Example 3:
Input: [1,3,5,6], 7
Output: 4

Example 4:
Input: [1,3,5,6], 0
Output: 0

予め昇順にソートされたリストnumsに対して、ある値targetが含まれていればそのインデックスを、含まれていなければリストの値が昇順に並ぶようにtargetがインサートされるべきインデックスを返す問題です。

解答・解説

解法1

まずはシンプルにiterativeに解決する方法を考えます。
numsにtargetが含まれていれば簡単ですから、それ以外のケースについて考えます。
numsにtargetが含まれていない場合ですが、2通りのケースがあります。
1. numsにtargetより大きな値が含まれる
2. numsにtargetより大きな値が含まれない
1のケースでは、はじめてtargetより大きな値が現れたとき、その値のインデックスにtargetを挿入すればよいわけなので、numsにtargetが含まれている場合と同じ処理で問題ありません。

if nums[i] >= tareget:
    return i

2のケースでは、numsの末尾にtargetを挿入することになるので、len(nums)がtargetを挿入すべきインデックスになります。ループ処理がリストの終点まで到達したときに、len(nums)を返すようにします。

elif i == len(nums) - 1:
    return len(nums)

まとめると、以下のコードになります。

class Solution:
    def searchInsert(self, nums: List[int], target: int) -> int:
        for i,n in enumerate(nums):
            if nums[i] >= target:
                return i
            elif i == len(nums) - 1:
                return len(nums)

解法2

Iterativeな解法だと時間計算量は $O(n)$ になってしまいますが、これを改善する二分探索（バイナリーサーチ）という方法があります。
以下記事が詳しいですが、
相手に「0 以上 32 未満の整数値」をなにか 1 つ思い浮かべてもらいその値を当てるとき、 0ですか？ 1ですか？ 2ですか？と順に聞いていくよりも
16以上ですか？ ->(yes)-> 24以上ですか？ ->(yes)-> 28以上ですか？と聞いていく方が早い
という発想です。

qiita.com

コードは以下の通りです。
探索範囲の始点と終点のインデックスを格納する変数としてleft, rightを用意し、その平均midのインデックスにあるnums[mid]がtarget以上かどうかで、探索範囲を徐々に狭めていきます。left, rightの大小関係がひっくり返った時点で終了します。

class Solution:
    def searchInsert(self, nums: List[int], target: int) -> int:
        left, right = 0, len(nums) - 1
        while left <= right:
            mid = int(left + (right - left) / 2)
            if nums[mid] >= target:
                right = mid - 1
            else:
                left = mid + 1
        return left

これで、時間計算量は $O(logn)$ に抑えられます。
ちなみに、中間の値を計算する際にmid = int(left + (right - left) / 2)としていますが、mid = int((left + right) / 2)ではなぜいけないのか、と思われた方もいらっしゃると思いますが、これはleft + rightの計算を先にしてしまうと値がオーバーフローしてしまう可能性があるための、ちょっとした細工です。