Редакция для Сумма квадратов на отрезке


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Автор: montes332

1. Идея

В каждом запросе нужно быстро находить сумму квадратов элементов на отрезке [l, r]:

a_l^2 + a_{l+1}^2 + ... + a_r^2

Если для каждого запроса считать эту сумму напрямую, то в худшем случае придется проходить почти весь массив, и это будет слишком медленно при n, q до 2 * 10^5.

Поэтому нужно сделать предварительную обработку массива: построить массив префиксных сумм квадратов. Тогда ответ на любой запрос будет находиться за O(1).


2. Наблюдения

Пусть pref[i] — сумма квадратов первых i элементов массива:

pref[i] = a_1^2 + a_2^2 + ... + a_i^2

Тогда:

  • pref[0] = 0
  • pref[1] = a_1^2
  • pref[2] = a_1^2 + a_2^2
  • и так далее

Чтобы получить сумму квадратов на отрезке [l, r], достаточно взять:

pref[r] - pref[l - 1]

Почему это работает:

  • pref[r] содержит сумму квадратов от 1 до r
  • pref[l - 1] содержит сумму квадратов от 1 до l - 1
  • если вычесть второе из первого, останется только отрезок от l до r

Еще одно важное замечание: квадраты могут быть большими. Например, 10000^2 = 100000000, а таких чисел много, поэтому нужно использовать 64-битный целый тип:

  • в C++ это long long
  • в Python обычный int уже подходит

3. Алгоритм

  1. Считать n и q.
  2. Считать массив из n чисел.
  3. Создать массив pref длины n + 1.
  4. Положить pref[0] = 0.
  5. Для каждого i от 1 до n вычислить:
    • pref[i] = pref[i - 1] + a[i - 1] * a[i - 1] в Python
    • или аналогично в C++
  6. Для каждого запроса считать l и r.
  7. Вывести pref[r] - pref[l - 1].

4. Почему это работает

Докажем корректность.

Построим массив pref, где pref[i] — сумма квадратов первых i элементов. Тогда по определению:

pref[r] = a_1^2 + a_2^2 + ... + a_{l-1}^2 + a_l^2 + ... + a_r^2

А также:

pref[l - 1] = a_1^2 + a_2^2 + ... + a_{l-1}^2

Если вычесть вторую сумму из первой, все слагаемые от 1 до l - 1 сократятся, и останется:

a_l^2 + a_{l+1}^2 + ... + a_r^2

Это и есть ответ на запрос.

Значит, формула pref[r] - pref[l - 1] действительно правильно находит сумму квадратов на отрезке [l, r].


5. Сложность

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

  • O(n)

Обработка каждого запроса:

  • O(1)

Общая сложность:

  • O(n + q)

Дополнительная память:

  • O(n)

6. Код на C++17

#include <iostream>
#include <vector>

using namespace std;

int main() {
    int n, q;
    cin >> n >> q;

    vector<long long> pref(n + 1, 0);

    for (int i = 1; i <= n; i++) {
        long long x;
        cin >> x;
        pref[i] = pref[i - 1] + x * x;
    }

    for (int i = 0; i < q; i++) {
        int l, r;
        cin >> l >> r;
        cout << pref[r] - pref[l - 1] << '\n';
    }

    return 0;
}

7. Код на Python 3

n, q = map(int, input().split())
a = list(map(int, input().split()))

pref = [0] * (n + 1)
for i in range(1, n + 1):
    pref[i] = pref[i - 1] + a[i - 1] * a[i - 1]

for _ in range(q):
    l, r = map(int, input().split())
    print(pref[r] - pref[l - 1])

Комментарии

Еще нет ни одного комментария.