Нахождение вставляемого элемента в отсортированном массиве

Дан отсортированный массив целых чисел nums и целое число target. Напишите функцию, которая находит индекс, на который можно вставить target, чтобы сохранить порядок массива. Вставляемый элемент может быть равен уже существующим элементам массива.

Пример:

  1. Вход: nums = [1, 3, 5, 6], target = 5 Выход: 2
    Объяснение: Число 5 уже присутствует в массиве на индексе 2.
  2. Вход: nums = [1, 3, 5, 6], target = 2 Выход: 1
    Объяснение: Число 2 должно быть вставлено на индекс 1, чтобы сохранить порядок.
  3. Вход: nums = [1, 3, 5, 6], target = 7 Выход: 4
    Объяснение: Число 7 должно быть вставлено на индекс 4, чтобы сохранить порядок.
  4. Вход: nums = [1, 3, 5, 6], target = 0 Выход: 0
    Объяснение: Число 0 должно быть вставлено на индекс 0, чтобы сохранить порядок.

Мое решение

Я исходил из того принципа, что используя бинарный поиск я либо найду индекс существующего в массиве элемента или не буду попадать в диапазоны его частей. При этом я решил, что есть два пограничных случая, которые буду проверять в начале функции - что у нас target находится вообще вне диапазона исходного массива, то есть он либо меньше самого маленького, либо больше самого большого.

Для подсказки себе в будущем я написал небольшой тестовый класс, в который добавил все примеры

@Test  
fun makeTest() {  
    Assert.assertEquals(2, searchInsert(intArrayOf(1, 3, 5, 6), 5))  
    Assert.assertEquals(1, searchInsert(intArrayOf(1, 3, 5, 6), 2))  
    Assert.assertEquals(4, searchInsert(intArrayOf(1, 3, 5, 6), 7))  
    Assert.assertEquals(0, searchInsert(intArrayOf(1, 3, 5, 6), 0))  
}

И само решени

fun searchInsert(nums: IntArray, target: Int): Int {  
    var start = 0  
    var end = nums.size - 1  
  
    if (nums[0] > target) return 0  
    if (nums[end] < target) return end + 1  
  
    while (start <= end) {  
        val midi = start + (end - start) / 2  
        val midva = nums[midi]  
  
        when {  
            midva == target -> return midi  
            target > nums[start] && target < midva -> end = midi - 1  
            target > midva && target < nums[end] -> start = midi + 1  
            else -> return midi + 1  
        }  
    }  
  
    return -1  
}

Компилятор и дебаггеры хороши тем, что автоматически и быстро пробегают по циклу и определяют корректность теста, но очень неплохо самому это делать. Чтобы не запутаться в вычислениях, "решаем" цикл чуть ниже. И вот что получается

Первый пример

target=5  
start=0, end=3, midi=1, midva=3 -> !if1, !if2, if3 -> start=2  
start=2, end=3, midi=2, midva=5 -> if1 -> return 2

Второй пример

target=2  
start=0, end=3, midi=1, midva=3 -> !if1, if2 -> end=0  
start=0, end=0, midi=0, midva=1 -> !if1, !if2, !if3 -> return 1

И так далее.

Исправленное решение

Ваше решение использует бинарный поиск, чтобы найти индекс для вставки целевого значения target в отсортированный массив nums. Рассмотрим подробно, как ваше решение работает и какие есть замечания:

Принцип работы вашего кода

  1. Инициализация границ:
    • start устанавливается в 0.
    • end устанавливается в nums.size - 1.
  2. Проверка крайних случаев:
    • Если target меньше первого элемента, возвращаем 0 (вставка в начало).
    • Если target больше последнего элемента, возвращаем end + 1 (вставка в конец).
  3. Бинарный поиск:
    • Находим середину midi и значение midva в этом индексе.
    • Используем конструкцию when для определения следующего шага:
      • Если midva равен target, возвращаем midi.
      • Если target больше nums[start] и меньше midva, сужаем поиск влево (end = midi - 1).
      • Если target больше midva и меньше nums[end], сужаем поиск вправо (start = midi + 1).
      • Если ни одно из условий не выполнено, предполагаем, что target должен быть вставлен после midi, и возвращаем midi + 1.
  4. Возврат значения:
    • Если не удалось найти индекс, возвращаем -1 (хотя это условие никогда не срабатывает, если входные данные корректны).

Замечания и исправления

  • Граничные случаи: Проверки на крайние случаи (nums[0] > target и nums[end] < target) правильные и предотвращают ненужные итерации, но в условии else нужно корректно обрабатывать все возможные случаи, когда target может быть вставлен в середину.
  • Поиск в else условии: В условии else предполагается, что target вставляется после текущего midi, но это не всегда правильно, если target равен midva. Вам нужно пересчитать условия для else.
  • Рассмотрение start == end: Этот случай не рассматривается отдельно, хотя его тоже нужно учесть, если target равен элементу в midi.

Исправленный код

fun searchInsert(nums: IntArray, target: Int): Int {
    var start = 0
    var end = nums.size - 1

    // Проверка крайних случаев
    if (target <= nums[start]) return start
    if (target > nums[end]) return end + 1

    while (start <= end) {
        val mid = start + (end - start) / 2
        val midVal = nums[mid]

        when {
            midVal == target -> return mid
            target < midVal -> end = mid - 1
            else -> start = mid + 1
        }
    }

    // Позиция, куда можно вставить target
    return start
}

Объяснение исправлений:

  1. Проверка крайних случаев: Обрабатываются правильно. Если target меньше или равен первому элементу, возвращается start, и если target больше последнего элемента, возвращается end + 1.

  2. Исправлен бинарный поиск:

    • Если midVal равен target, возвращаем mid.
    • Если target меньше midVal, сужаем диапазон поиска влево.
    • Если target больше midVal, сужаем диапазон поиска вправо.

    Эти изменения гарантируют, что все возможные случаи корректно обрабатываются, и start будет указывать на индекс, куда можно вставить target, когда поиск завершен.

  3. Возврат значения: В конце, возвращаем start, который указывает на правильную позицию для вставки target.