Редакция для Количество пар с суммой K в отсортированном массиве


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] + a[r]:

  • если сумма меньше K, нужно увеличить сумму, значит двигаем l вправо;
  • если сумма больше K, нужно уменьшить сумму, значит двигаем r влево;
  • если сумма равна K, нужно аккуратно посчитать, сколько пар дают текущие значения, особенно если есть одинаковые элементы.

Главная сложность задачи — корректно обработать повторяющиеся числа.


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

Наблюдение 1. Почему работают два указателя

Так как массив неубывающий:

  • если a[l] + a[r] < K, то с текущим l никакой индекс левее r не поможет получить K, потому что все такие элементы не больше a[r];
  • значит, текущий l можно отбросить и увеличить l.

Аналогично:

  • если a[l] + a[r] > K, то с текущим r никакой индекс правее l не поможет получить K, потому что все такие элементы не меньше a[l];
  • значит, текущий r можно отбросить и уменьшить r.

Это позволяет пройти массив за линейное время.

Наблюдение 2. Если найдено совпадение, одинаковых элементов может быть много

Пусть a[l] + a[r] == K.

Тогда есть два случая.

Случай 1: a[l] == a[r]

Раз массив отсортирован, это значит, что все элементы на отрезке от l до r одинаковые.

Если на этом отрезке cnt элементов, то нужно выбрать любые два разных индекса. Число таких пар равно:

cnt * (cnt - 1) / 2

После этого можно завершать алгоритм: внутри этого отрезка уже учтены все возможные пары.

Случай 2: a[l] != a[r]

Пусть:

  • слева значение leftValue = a[l] встречается cntL раз подряд,
  • справа значение rightValue = a[r] встречается cntR раз подряд.

Тогда каждая копия слева может образовать пару с каждой копией справа, поэтому добавляем:

cntL * cntR

После этого сдвигаем оба указателя за эти группы одинаковых элементов.


3. Алгоритм

  1. Считать n, K и массив a.
  2. Установить два указателя:
    • l = 0,
    • r = n - 1.
  3. Завести переменную ans = 0.
  4. Пока l < r:
    • вычислить sum = a[l] + a[r];
    • если sum < K, увеличить l;
    • если sum > K, уменьшить r;
    • иначе:
      • если a[l] == a[r]:
        • cnt = r - l + 1;
        • добавить к ответу cnt * (cnt - 1) / 2;
        • завершить цикл;
      • иначе:
        • посчитать, сколько одинаковых элементов подряд стоит слева: cntL;
        • посчитать, сколько одинаковых элементов подряд стоит справа: cntR;
        • добавить cntL * cntR к ответу.
  5. Вывести ans.

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

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

1. Если сумма меньше K, левый указатель можно сдвинуть

Пусть a[l] + a[r] < K.

Так как массив отсортирован, для любого j, где l < j <= r, выполнено a[j] <= a[r].

Значит: a[l] + a[j] <= a[l] + a[r] < K

То есть элемент a[l] вообще не может участвовать ни в одной подходящей паре с индексом справа от него в текущем диапазоне. Поэтому его можно безопасно отбросить, увеличив l.

2. Если сумма больше K, правый указатель можно сдвинуть

Пусть a[l] + a[r] > K.

Для любого i, где l <= i < r, выполнено a[i] >= a[l].

Значит: a[i] + a[r] >= a[l] + a[r] > K

То есть элемент a[r] не может участвовать ни в одной подходящей паре с индексом слева от него в текущем диапазоне. Поэтому его можно безопасно отбросить, уменьшив r.

3. Если сумма равна K и значения на концах разные

Пусть a[l] + a[r] == K, a[l] = x, a[r] = y, и x != y.

Пусть слева подряд стоит cntL элементов, равных x, а справа подряд стоит cntR элементов, равных y.

Каждый из cntL элементов со значением x образует подходящую пару с каждым из cntR элементов со значением y, потому что x + y = K.

Иных пар с этими крайними значениями нет: все такие элементы уже учтены. Поэтому можно добавить cntL * cntR и перейти к следующим отличающимся значениям.

4. Если сумма равна K и значения на концах одинаковые

Пусть a[l] == a[r]. Так как массив отсортирован, все элементы между l и r равны одному и тому же значению.

Если их количество равно cnt, то любая пара разных индексов на этом отрезке даёт сумму K. Число таких пар — это количество способов выбрать 2 элемента из cnt, то есть cnt * (cnt - 1) / 2.

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

5. Каждая подходящая пара учитывается ровно один раз

Алгоритм либо отбрасывает элементы, которые точно не могут участвовать в ответе, либо сразу считает все пары для крайних групп одинаковых значений и сдвигает указатели дальше.

Из-за этого:

  • ни одна подходящая пара не пропускается;
  • ни одна пара не считается дважды.

Значит, итоговый ответ верен.


5. Сложность

Каждый указатель движется только в одну сторону, и каждый элемент массива обрабатывается не более одного раза.

  • Время: O(n)
  • Память: O(1), не считая массива

6. Код на C++17

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int n;
    long long K;
    cin >> n >> K;

    vector<long long> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }

    int l = 0, r = n - 1;
    long long ans = 0;

    while (l < r) {
        long long sum = a[l] + a[r];

        if (sum < K) {
            l++;
        } else if (sum > K) {
            r--;
        } else {
            if (a[l] == a[r]) {
                long long cnt = r - l + 1;
                ans += cnt * (cnt - 1) / 2;
                break;
            } else {
                long long leftValue = a[l];
                long long rightValue = a[r];
                long long cntL = 0, cntR = 0;

                while (l <= r && a[l] == leftValue) {
                    cntL++;
                    l++;
                }
                while (l <= r && a[r] == rightValue) {
                    cntR++;
                    r--;
                }

                ans += cntL * cntR;
            }
        }
    }

    cout << ans << '\n';
    return 0;
}

7. Код на Python 3

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

l = 0
r = n - 1
ans = 0

while l < r:
    s = a[l] + a[r]

    if s < K:
        l += 1
    elif s > K:
        r -= 1
    else:
        if a[l] == a[r]:
            cnt = r - l + 1
            ans += cnt * (cnt - 1) // 2
            break
        else:
            left_value = a[l]
            right_value = a[r]
            cnt_l = 0
            cnt_r = 0

            while l <= r and a[l] == left_value:
                cnt_l += 1
                l += 1

            while l <= r and a[r] == right_value:
                cnt_r += 1
                r -= 1

            ans += cnt_l * cnt_r

print(ans)

Комментарии

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