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

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

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

AtCoder AtCoder400点 ABC-D 差分更新 二項係数 数え上げ問題 集計処理 各kに対して バケット nC2 解空間:O(N^2)通りの選択肢 茶色diff 典型要素を詰め合わせた教育的問題 N個のものうち1個を変更・削除したものを解く 連想配列(setやmap)

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

問題へのリンク

問題概要

長さ 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();
}