Индекс равновесия

Постановка задачи

nums = [1, 7, 3, 6, 5, 6]

3

nums = [1, 2, 3]

-1

nums = [2, 1, -1]

0

Решение

Эта задача хорошо показывает, зачем нужны префиксные суммы. Для каждого индекса нам нужно знать две величины:

1. сумму всех элементов слева;
2. сумму всех элементов справа.

Наивное решение для каждого индекса заново считает сумму слева и сумму справа. Такой подход легко реализовать, но он делает много лишней работы. Для массива длины n мы будем пересчитывать похожие суммы снова и снова, поэтому получим сложность O(n²).

Нам нужно заметить, что при движении слева направо сумма слева меняется очень просто: после проверки текущего индекса мы добавляем в нее nums[i].

Идея решения:

Сначала посчитаем общую сумму массива total. Потом будем идти по массиву и поддерживать переменную leftSum - сумму элементов слева от текущего индекса.

Если мы знаем total, leftSum и текущий элемент nums[i], то сумму справа можно получить без отдельного прохода:

rightSum = total - leftSum - nums[i];

Теперь на каждом индексе достаточно проверить leftSum === rightSum. Если равенство выполняется, возвращаем текущий индекс. Так как мы идем слева направо, первый найденный индекс автоматически будет самым левым.

Основные шаги:

1. Считаем total - сумму всех элементов массива.

2. Создаем переменную leftSum = 0.

3. Идем по массиву слева направо:

   • считаем rightSum = total - leftSum - nums[i];
   • если leftSum == rightSum, возвращаем i;
   • иначе добавляем текущий элемент к leftSum.

4. Если прошли весь массив и ничего не нашли, возвращаем -1.

Рассмотрим это на примере:

nums = [1, 7, 3, 6, 5, 6]

Общая сумма массива:

total = 28

Идем по индексам:

1. Индекс 0, число 1.

   • leftSum = 0
   • rightSum = 28 - 0 - 1 = 27
   • суммы не равны, добавляем 1 к leftSum

2. Индекс 1, число 7.

   • leftSum = 1
   • rightSum = 28 - 1 - 7 = 20
   • суммы не равны, добавляем 7 к leftSum

3. Индекс 2, число 3.

   • leftSum = 8
   • rightSum = 28 - 8 - 3 = 17
   • суммы не равны, добавляем 3 к leftSum

4. Индекс 3, число 6.

   • leftSum = 11
   • rightSum = 28 - 11 - 6 = 11
   • суммы равны, возвращаем 3

Тут важно не добавлять текущий элемент к leftSum до проверки. По условию индекс равновесия сравнивает элементы строго слева и строго справа, а сам nums[i] не входит ни в одну из этих сумм.

Оценка сложности

Временная сложность

Мы один раз считаем общую сумму массива и один раз проходим по массиву в поиске индекса равновесия. Получаем O(n).

Пространственная сложность

Мы используем только несколько переменных: total, leftSum и rightSum. Дополнительная память не зависит от размера массива, поэтому пространственная сложность равна O(1).

Код решения

Приведем код решения.

func pivotIndex(nums []int) int {
    total := 0
    for _, num := range nums {
        total += num
    }
 
    leftSum := 0
    for i, num := range nums {
        rightSum := total - leftSum - num
        if leftSum == rightSum {
            return i
        }
        leftSum += num
    }
    return -1
}

Итоги

Задача учит не пересчитывать суммы с нуля.

1. Ключевая идея. Если известна общая сумма массива, сумму справа можно получить через total - leftSum - nums[i].

2. Порядок действий. Сначала сравниваем leftSum и rightSum, и только потом добавляем текущий элемент к leftSum.

3. Почему возвращаем первый найденный индекс. Мы идем слева направо, поэтому первый подходящий индекс будет самым левым.

4. Оптимизация памяти. Для этой задачи не нужен отдельный массив префиксных сумм, достаточно одной переменной.