Редакция для Количество пар в двух отсортированных массивах с суммой K
Submitting an official solution before solving the problem yourself is a bannable offence.
Автор:
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. Алгоритм
- Считываем
n,m,K, затем массивыaиb. - Инициализируем:
i = 0,j = m - 1,ans = 0.
- Пока
i < nиj >= 0:- вычисляем
sum = a[i] + b[j]; - если
sum < K, увеличиваемi; - если
sum > K, уменьшаемj; - если
sum == K:- запоминаем значения
va = a[i]иvb = b[j]; - считаем, сколько раз подряд
vaвстречается вa, начиная с позицииi; - считаем, сколько раз подряд
vbвстречается вb, начиная с позицииjвлево; - добавляем к ответу
cntA * cntB.
- запоминаем значения
- вычисляем
- Выводим
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)
Комментарии