Сумма диапазона в неизменяемом массиве

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

sumRange(left, right);

Input:
["NumArray", "sumRange", "sumRange", "sumRange"]
[[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]]

[null, 1, -1, -3]

const numArray = new NumArray([-2, 0, 3, -5, 2, -1]);
numArray.sumRange(0, 2); // return 1  (-2 + 0 + 3)
numArray.sumRange(2, 5); // return -1 (3 + -5 + 2 + -1)
numArray.sumRange(0, 5); // return -3 (-2 + 0 + 3 + -5 + 2 + -1)

Решение

Условие почти напрямую указывает на использование префиксных сумм. Нам нужно много раз отвечать на запрос: "какая сумма элементов на отрезке от left до right?"

Если каждый раз проходить по этому отрезку циклом, один запрос будет работать за O(k), где k - длина диапазона. В худшем случае это O(n) на один вызов sumRange. При 10⁴ запросах такая повторная работа быстро становится неэффективной.

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

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

Создадим массив prefix, где prefix[i] хранит сумму элементов от начала массива до индекса i включительно. То есть:

prefix[0] = nums[0];
prefix[1] = nums[0] + nums[1];
prefix[2] = nums[0] + nums[1] + nums[2];
prefix[3] = nums[0] + nums[1] + nums[2] + nums[3];

Для массива:

nums = [-2, 0, 3, -5, 2, -1];

получим:

prefix = [-2, -2, 1, -4, -2, -3];

Теперь сумма от left до right включительно считается так:

left === 0 ? prefix[right] : prefix[right] - prefix[left - 1];

В prefix[right] лежит сумма всех элементов от начала массива до right. Когда left > 0, в prefix[left - 1] лежит сумма всех элементов до left, не включая left. Если вычесть одно из другого, останется ровно нужный диапазон.

Если left = 0, вычитать нечего: нужная сумма уже лежит в prefix[right].

Например, для запроса:

sumRange(2, 5);

нам нужна сумма:

nums[2] + nums[3] + nums[4] + nums[5];

По массиву префиксных сумм это:

prefix[5] - prefix[1] = -3 - (-2) = -1

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

1. В конструкторе создаем массив prefix длины nums.length.

2. Записываем nums[0] в prefix[0], а дальше заполняем массив так, чтобы prefix[i] = prefix[i - 1] + nums[i].

3. В sumRange(left, right) возвращаем prefix[right], если left = 0, иначе prefix[right] - prefix[left - 1].

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

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

Конструктор один раз проходит по массиву и строит prefix, поэтому работает за O(n).

Каждый вызов sumRange делает несколько простых операций и работает за O(1).

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

Мы храним массив prefix длины n, поэтому используем O(n) дополнительной памяти.

Код решения

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

type NumArray struct {
    prefix []int
}
 
func Constructor(nums []int) NumArray {
    prefix := make([]int, len(nums))
    prefix[0] = nums[0]
    for i := 1; i < len(nums); i++ {
        prefix[i] = prefix[i-1] + nums[i]
    }
    return NumArray{prefix: prefix}
}
 
func (this *NumArray) SumRange(left int, right int) int {
    if left == 0 {
        return this.prefix[right]
    }
    return this.prefix[right] - this.prefix[left-1]
}

Итоги

Задача иллюстрирует классический случай, когда стоит использовать O(n) дополнительной памяти и времени на подготовку префиксной суммы, чтобы потом отвечать на каждый запрос за O(1).

1. Ключевая идея. prefix[i] хранит сумму от начала массива до индекса i включительно.

2. Формула запроса. Сумма от left до right равна prefix[right], если left = 0, иначе prefix[right] - prefix[left - 1].

3. Почему это работает. Если left > 0, из суммы до правой границы мы вычитаем сумму до левой границы, поэтому остается только нужный диапазон.

4. Важно, что массив неизменяемый. Мы строим prefix один раз и больше не пересчитываем его.