abc153_dの解き方の解説 Banned K nCrをメモ化しておくのがコツ

2020/05/23 05:05

AtCoder Beginner Contest 153のD問題 Bouquetを解く。必要となる素数で割ったあまりを取得するクラスと組み合わせ数の計算方法の解決

感謝の正拳突き１万回ならぬ、競プロサイトで 1 日 1AC(正解)するように頑張る。

今日はAtCoder Beginner Contest 159 の D 問題 Banned Kです。

Banned って垢 BAN とかで使う BAN の過去形か・・

考え方

数字別に個数を集計する
個数ごとに nC2 で組み合わせ数を計算して、どれも BAN されていないサマリーする
サマリー - Ai の組み合わせ数 + Ai を 1 つ取り除いた組み合わせ数 = 答え

ポイント

N の制約が 200,000 個以下と比較的緩めです。ゴリゴリ計算しても間に合う可能性があります。

とはいえ何も考えずに組み合わせ数を毎回計算してしまうと、TLE してしまいます。1 回計算した値をメモしておくメモ化は必須となります。

実装はabc156_d の解説で組み合わせ数 nCr をメモ化しながら計算する関数を紹介しているので、そちらを参照してください。

#include <bits/stdc++.h>

using namespace std;
#define ll long long

// 組み合わせ数を計算する。nCr
unordered_map<ll, vector<ll>> _combinations;
ll combination(int n, int r) {
    if(n < 1 || r < 1 || n < r) return 0;
    if(n == r) return 1;
    if(r > n - r) r = n - r;
    if(!_combinations.count(n)) {
        vector<ll> vec = {0, n};
        _combinations.emplace_hint(_combinations.end(), n, vec);
    }
    auto&& vec = _combinations[n];
    if(r > n - r) r = n - r;
    if(vec.size() <= r) {
        int i = vec.size();
        vec.resize(r + 1);
        for(; i <= r; i++) {
            vec[i] = vec[i - 1];
            vec[i] *= (n - i + 1);
            vec[i] /= i;
        }
    }

    return vec[r];
}
int main() {
    cin.tie(0); ios::sync_with_stdio(false); cout << fixed << std::setprecision(20);

    int n;
    cin >> n;
    vector<ll> as(n), counts(n + 1, 0);

    // 1. 数字別に個数を集計する
    for(auto&& a : as) {
        cin >> a;
        counts[a]++;
    }

    // 2. 個数ごとにnC2で組み合わせ数を計算してサマリーする
    ll total = 0;
    for(auto&& count : counts) {
        if(count > 0) {
            total += combination(count, 2);
        }
    }

    for(auto&& a : as) {
        // 3. サマリー - Aiの組み合わせ数 + Aiを1つ取り除いた組み合わせ数 = 答え
        cout << total - combination(counts[a], 2) + combination(counts[a] - 1, 2) << endl;
    }
}