Редакция для Количество пар в двух отсортированных массивах с суммой 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. Идея

Даны два массива a и b, оба уже отсортированы по неубыванию. Нужно посчитать число пар индексов (i, j), для которых a[i] + b[j] = K.

Перебирать все пары нельзя: это заняло бы O(n * m), что слишком долго при n, m до 2 * 10^5.

Ключевая идея — использовать метод двух указателей:

  • указатель i ставим в начало массива a,
  • указатель j ставим в конец массива b.

Тогда:

  • если a[i] + b[j] < K, нужно увеличить сумму, значит двигаем i вправо;
  • если a[i] + b[j] > K, нужно уменьшить сумму, значит двигаем j влево;
  • если a[i] + b[j] = K, нужно аккуратно учесть все одинаковые значения подряд.

Именно последний пункт особенно важен: если, например, a[i] встречается несколько раз подряд, и b[j] тоже встречается несколько раз подряд, то каждая копия из первого массива образует пару с каждой копией из второго.


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

Наблюдение 1. Отсортированность позволяет двигаться только вперёд

Если сейчас a[i] + b[j] < K, то для текущего i и любого индекса левее j сумма будет не больше, потому что b отсортирован. Значит, уменьшать j бессмысленно. Единственный шанс получить K — увеличить a[i], то есть сдвинуть i.

Аналогично, если a[i] + b[j] > K, то увеличивать i бессмысленно, потому что сумма станет ещё больше. Нужно уменьшать j.

Наблюдение 2. Одинаковые элементы нужно считать группами

Пусть найдено, что a[i] + b[j] = K.

Если в массиве a подряд идут cntA одинаковых значений va, а в массиве b подряд идут cntB одинаковых значений vb, и va + vb = K, то число подходящих пар равно cntA * cntB.

Почему? Потому что любой из этих cntA элементов первого массива можно объединить с любым из этих cntB элементов второго массива.

Наблюдение 3. Каждый элемент будет просмотрен не более одного раза

Указатель i только увеличивается, указатель j только уменьшается. Значит, общий проход линейный.


3. Алгоритм

  1. Считываем n, m, K, затем массивы a и b.
  2. Инициализируем:
    • i = 0,
    • j = m - 1,
    • ans = 0.
  3. Пока i < n и j >= 0:
    • вычисляем sum = a[i] + b[j];
    • если sum < K, увеличиваем i;
    • если sum > K, уменьшаем j;
    • если sum == K:
      1. запоминаем значения va = a[i] и vb = b[j];
      2. считаем, сколько раз подряд va встречается в a, начиная с позиции i;
      3. считаем, сколько раз подряд vb встречается в b, начиная с позиции j влево;
      4. добавляем к ответу cntA * cntB.
  4. Выводим ans.

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

Докажем, что алгоритм считает ответ правильно.

Случай 1. a[i] + b[j] < K

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

  • a[i] + b[j'] <= a[i] + b[j] < K.

Значит, с текущим a[i] нельзя получить сумму K ни с одним элементом b на позициях от 0 до j. Следовательно, элемент a[i] можно безопасно отбросить, увеличив i.

Случай 2. a[i] + b[j] > K

Так как массив a отсортирован, для любого i' >= i выполняется a[i'] >= a[i]. Тогда:

  • a[i'] + b[j] >= a[i] + b[j] > K.

Значит, с текущим b[j] нельзя получить сумму K ни с одним элементом a на позициях от i до n - 1. Следовательно, элемент b[j] можно безопасно отбросить, уменьшив j.

Случай 3. a[i] + b[j] = K

Пусть:

  • в a начиная с позиции i стоит cntA одинаковых элементов va,
  • в b начиная с позиции j влево стоит cntB одинаковых элементов vb.

Так как va + vb = K, каждая из cntA позиций первого массива подходит к каждой из cntB позиций второго. Значит, число новых пар равно cntA * cntB.

После этого все эти элементы можно пропустить:

  • другие пары с элементом va и значением b, отличным от vb, уже не дадут нужную сумму;
  • другие пары с элементом vb и значением a, отличным от va, тоже не дадут нужную сумму в текущем диапазоне поиска.

Поэтому мы не пропускаем ни одной подходящей пары и не считаем ни одну пару дважды.

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


5. Сложность

Каждый указатель проходит свой массив не более одного раза:

  • i увеличивается от 0 до n,
  • j уменьшается от m - 1 до -1.

Поэтому:

  • время работы: O(n + m),
  • дополнительная память: O(1), не считая хранения входных массивов.

6. Код на C++17

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

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

    vector<long long> a(n), b(m);
    for (int i = 0; i < n; i++) cin >> a[i];
    for (int i = 0; i < m; i++) cin >> b[i];

    int i = 0;
    int j = m - 1;
    long long ans = 0;

    while (i < n && j >= 0) {
        long long sum = a[i] + b[j];

        if (sum < K) {
            i++;
        } else if (sum > K) {
            j--;
        } else {
            long long va = a[i];
            long long vb = b[j];

            long long cntA = 0;
            while (i < n && a[i] == va) {
                cntA++;
                i++;
            }

            long long cntB = 0;
            while (j >= 0 && b[j] == vb) {
                cntB++;
                j--;
            }

            ans += cntA * cntB;
        }
    }

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

7. Код на Python 3

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

i = 0
j = m - 1
ans = 0

while i < n and j >= 0:
    s = a[i] + b[j]

    if s < K:
        i += 1
    elif s > K:
        j -= 1
    else:
        va = a[i]
        vb = b[j]

        cnt_a = 0
        while i < n and a[i] == va:
            cnt_a += 1
            i += 1

        cnt_b = 0
        while j >= 0 and b[j] == vb:
            cnt_b += 1
            j -= 1

        ans += cnt_a * cnt_b

print(ans)

Комментарии

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