Редакция для Пары равных элементов


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.

1. Идея

Нужно посчитать количество пар индексов (i, j), где i < j и a[i] = a[j].

Прямой перебор всех пар слишком медленный: пар порядка n^2, а n может быть до 2 * 10^5.

Главная идея — обрабатывать элементы слева направо и хранить, сколько раз каждое число уже встречалось.

Когда мы читаем очередное значение x:

  • если раньше x встречалось k раз,
  • то с текущим элементом оно образует ровно k новых пар.

Значит, достаточно:

  1. узнать текущую частоту x,
  2. прибавить её к ответу,
  3. увеличить частоту x на 1.

Для хранения частот удобно использовать хеш-таблицу: unordered_map в C++ или dict в Python.


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

Наблюдение 1

Если число x встречается k раз, то среди этих вхождений можно составить количество пар:

k * (k - 1) / 2

Но не обязательно сначала считать все частоты, а потом отдельно суммировать. Можно делать это сразу по ходу чтения массива.

Наблюдение 2

Пусть мы уже обработали часть массива. Если очередной элемент равен x, а x уже встречался freq[x] раз, то:

  • каждая из этих прошлых позиций вместе с текущей образует одну новую пару;
  • значит, в ответ нужно добавить freq[x].

После этого текущий элемент тоже становится одним из уже встреченных, поэтому нужно сделать freq[x] += 1.

Наблюдение 3

Ответ может быть большим. Например, если все n чисел одинаковы, то ответ равен числу пар среди n элементов. Поэтому в C++ нужен тип long long.


3. Алгоритм

  1. Считать n.
  2. Создать словарь частот freq.
  3. Завести переменную answer = 0.
  4. Для каждого элемента x:
    • прибавить к answer значение freq[x] или 0, если такого ключа ещё нет;
    • увеличить freq[x] на 1.
  5. Вывести answer.

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

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

Рассмотрим момент, когда обрабатывается элемент a[j] = x.

Нас интересуют все пары (i, j), где i < j и a[i] = a[j]. К этому моменту все элементы с индексами меньше j уже обработаны. Значит:

  • freq[x] хранит количество таких индексов i, что i < j и a[i] = x;
  • каждый такой индекс образует с j одну подходящую пару;
  • других подходящих пар с правым концом в j не существует.

Следовательно, число новых пар, которые появляются при обработке позиции j, равно ровно freq[x].

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

Таким образом:

  • каждая подходящая пара учитывается ровно один раз — когда обрабатывается её правая позиция;
  • никакая лишняя пара не добавляется.

Значит, итоговое значение answer равно количеству всех пар одинаковых элементов.


5. Сложность

Для каждого из n элементов выполняется константное число операций со словарём.

  • Время: O(n) в среднем.
  • Память: O(k), где k — число различных значений, в худшем случае O(n).

6. Код на C++17

#include <iostream>
#include <vector>
#include <unordered_map>

using namespace std;

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

    unordered_map<long long, long long> freq;
    long long answer = 0;

    for (int i = 0; i < n; i++) {
        long long x;
        cin >> x;
        answer += freq[x];
        freq[x]++;
    }

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

7. Код на Python 3

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

freq = {}
answer = 0

for x in a:
    answer += freq.get(x, 0)
    freq[x] = freq.get(x, 0) + 1

print(answer)

Комментарии

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