Произведение массива

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

nums = [1, 2, 3, 4]

[24, 12, 8, 6]

nums = [-1, 1, 0, -3, 3]

[0, 0, 9, 0, 0]

Решение

Это задача среднего уровня на LeetCode. Если вы знакомы с паттерном префиксных сумм, трудностей не возникнет.

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

Значит, нужно иначе получить произведение всех элементов, кроме текущего.

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

answer[i] = произведение слева от i * произведение справа от i

Текущий элемент nums[i] в это произведение не входит.

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

Сначала пройдем слева направо и для каждого индекса сохраним произведение всех элементов слева от него. Потом пройдем справа налево и умножим уже записанное значение на произведение всех элементов справа.

Для массива:

nums = [1, 2, 3, 4]

после первого прохода получим:

answer = [1, 1, 2, 6]

Здесь answer[i] хранит произведение элементов строго слева от i:

• слева от 1 ничего нет, поэтому произведение равно 1;
• слева от 2 стоит 1;
• слева от 3 стоит 1 * 2 = 2;
• слева от 4 стоит 1 * 2 * 3 = 6.

Теперь идем справа налево и поддерживаем переменную rightProduct - произведение всех элементов справа от текущего индекса.

На каждом шаге умножаем:

answer[i] = answer[i] * rightProduct;

После этого обновляем rightProduct, добавляя в него текущий элемент:

rightProduct = rightProduct * nums[i];

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

1. Создаем массив answer длины n.

2. Делаем проход слева направо:

   • поддерживаем leftProduct;
   • записываем в answer[i] произведение всех элементов слева;
   • после записи умножаем leftProduct на nums[i].

3. Делаем проход справа налево:

   • поддерживаем rightProduct;
   • умножаем answer[i] на произведение всех элементов справа;
   • после этого умножаем rightProduct на nums[i].

4. Возвращаем answer.

Разберем второй проход на примере:

nums = [1, 2, 3, 4]
answer после первого прохода = [1, 1, 2, 6]

Идем справа налево:

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

   • rightProduct = 1
   • answer[3] = 6 * 1 = 6
   • обновляем rightProduct = 1 * 4 = 4

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

   • rightProduct = 4
   • answer[2] = 2 * 4 = 8
   • обновляем rightProduct = 4 * 3 = 12

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

   • rightProduct = 12
   • answer[1] = 1 * 12 = 12
   • обновляем rightProduct = 12 * 2 = 24

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

   • rightProduct = 24
   • answer[0] = 1 * 24 = 24
   • обновляем rightProduct = 24 * 1 = 24

В итоге получаем:

answer = [24, 12, 8, 6]

Важно, что мы сначала используем текущее значение leftProduct или rightProduct, и только потом умножаем его на nums[i]. Так текущий элемент не попадает в собственный ответ.

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

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

Мы делаем два прохода по массиву: один слева направо и один справа налево. Получаем O(n).

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

Входной массив и массив ответа занимают O(n) памяти. Дополнительно наше решение использует только две переменные: leftProduct и rightProduct. Поэтому дополнительная память равна O(1).

Код решения

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

func productExceptSelf(nums []int) []int {
    n := len(nums)
    answer := make([]int, n)
    leftProduct := 1
    for i := 0; i < n; i++ {
        answer[i] = leftProduct
        leftProduct *= nums[i]
    }
 
    rightProduct := 1
    for i := n - 1; i >= 0; i-- {
        answer[i] *= rightProduct
        rightProduct *= nums[i]
    }
    return answer
}

Итоги

Задача учит применять идею префиксных значений не только к суммам, но и к произведениям.

1. Ключевая идея. Ответ для индекса i равен произведению всех элементов слева от i на произведение всех элементов справа от i.

2. Почему не нужно деление. Два прохода дают нам левую и правую части ответа напрямую, без деления и специальных условий для нулей.

3. Порядок обновления важен. Сначала используем накопленное произведение, и только потом умножаем его на текущий элемент.