差分更新系の教育的問題かな

問題概要

長さ $N$ の数列が与えらえる。各 index k に対して、

数列から k 番目の要素を除いたものについて
その中から異なる 2 つの要素を選ぶ方法であって (順番は問わず)
その 2 つの要素の選び方が等しい

という選び方が何通りあるのかを求めよ。

制約

$3 \le N \le 2 \times 10^{5}$

考えたこと

愚直に毎回の数列について計算したのでは $O(N^{2})$ の計算量を要してしまう。今回のような、各 $k$ について、それを除いたものについて考えるタイプの問題では

まず全体の数列についての答えを求めておいて
各要素を取り除くことによる「影響」を計算し、
その影響分を引く

という風にするのが定石。その影響を計算することをしばしば「差分更新」と呼んだりする。つまり、差分だけ高速に計算してしまおうという発想だ。

全体

何にしてもまずは、全体の数列について答えを求める必要がある。でもこれは簡単だ。

nums[ v ] := 値 v が何個あるか

という配列を用意してあげて、各値 v について、

$n$ = nums[ v ] として
$\frac{n(n-1)}{2}$ を加算する

という風にして求めることができる。

差分更新

ここまでわかれば差分計算も簡単だ。値 v を除くとき、影響があるのは nums[ v ] の部分だけだ。よって、その部分についてのみ、before / after を計算して補正すればよい。具体的には、

$n$ = nums[ v ] として
before： $\frac{n(n-1)}{2}$
after： $\frac{(n-1)(n-2)}{2}$

という風になる。答えは、(全体) - before + after となる。計算量は全体で $O(N)$ でできる。

#include <bits/stdc++.h>
using namespace std;

int N;
vector<long long> a;

void solve() {
    map<long long, long long> ma;
    for (int i = 0; i < N; ++i) ma[a[i]]++;
    long long res = 0;
    for (auto it : ma) res += it.second * (it.second - 1) / 2;
    for (int i = 0; i < N; ++i) {
        long long val = ma[a[i]];
        long long before = val * (val - 1) / 2;
        long long after = (val - 1) * (val - 2) / 2;
        cout << res + after - before << endl;
    }
}

int main() {
    cin >> N;
    a.resize(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    solve();
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 159 D - Banned K (3Q, 茶色, 400 点)

問題概要

制約

考えたこと

全体

差分更新