【AtCoder】ABC 392 D - Doubles | 緑コーダーが解くAtCoder

配点: 400 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 579

問題概要
制約
考察
実装例
コメント

問題概要

$N$ 個のサイコロがあり、 $i$ 番目のサイコロには $K_i$ 個の面があって、各面にはそれぞれ $A_{i,1},A_{i,2},\ldots,A_{i,K_i}$ が書かれている。このサイコロを振ると、各面がそれぞれ $\frac{1}{K_i}$ の確率で出る。

$N$ 個のサイコロから2個を選んで振るとき、2つのサイコロの出目が等しくなる確率の最大値を求めよ。

制約

$2 \leq N \leq 100$
$1 \leq K_i$
$K_1+K_2+\cdots+K_N \leq 10^5$
$1 \leq A_{i,j} \leq 10^5$
入力は全て整数。

考察

確率を考える上で必要なのは次の2つ。

ある事象 $A$ が起こる場合の数
全事象の場合の数

今回の場合は、

「 $N$ 個のサイコロから2個を選んで振ったとき、同じ出目が出る」という事象の場合の数
$N$ 個のサイコロから2個を選んで振ったとき、2つのサイコロの出目の組としてあり得る個数(ただし、同じ出目でも異なる面のものは区別して数える)

となる。

例えば、

$K_i$ 個の面を持ち、数 $X$ が書かれた面が $a$ 個
$K_j$ 個の面を持ち、数 $X$ が書かれた面が $b$ 個

という2つのサイコロを振ったとき、どちらも出目 $X$ が出る場合の数は $ab$ 通り存在する。

また、この2つのサイコロの出目の組の総数は $K_i K_j$ 通り存在するので、求める確率 $P_{i, j}(X)$ は、 $P_{i, j}(X) = \frac{ab}{K_i K_j}$ と計算できる。

それぞれの事象は排反であるから、面の種類ごとに和を取り、 $i, j$ の選び方をすべて試してそれを最大化すれば答えが求まる。

以下、具体的なアルゴリズムを考えていく。

まず、

$\mathrm{tables}_{i}(X) :=$ サイコロ $i$ の面の中で、数 $X$ が書かれているものの個数

という連想配列を用意する。その上で、以下のアルゴリズムを構築する。

$1 \leq i \lt j \leq N$ を満たす組み合わせを二重ループで探索
ある組み合わせについて、
1. カウンター $\mathrm{cnt}$ を $0$ で初期化する。
2. の各要素に対して、
  1. もし $\mathrm{tables}_{j}(X_k) = 0$ ならば、サイコロ $i, j$ に共通する出目は存在しないのでスキップ。
  2. そうでないとき、 $\mathrm{cnt}$ に $\mathrm{tables}_{i}(X_k) \times \mathrm{tables}_{j}(X_k)$ を加算する。
3. $p = \frac{\mathrm{cnt}}{K_i K_j}$ を計算。
それぞれの組み合わせごとに $p$ を求めて、これを最大化する。

計算量は、 $\displaystyle M = \sum K_i$ とれば、

$\mathrm{tables}_{i}(X)$ の構築 : $O(M \log M)$
$(i, j)$ の組み合わせの全探索 : $O(N)$

なので、全体としては $O(NM \log M)$ となり、これは本問の制約下では十分高速。

実装においては、 $\mathrm{cnt}$ 計算時のオーバーフローや、除算時のキャスト処理に注意しよう(3敗)。

実装例

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

#define fix(x) fixed << setprecision(x)
#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;
using lld = long double;

// ======================================== //

int main()
{
    int N;
    cin >> N;
    vector<ll> K(N);
    vector<map<int, ll>> tables(N);
    rep(i, 0, N)
    {
        cin >> K[i];
        rep(j, 0, K[i])
        {
            int A;
            cin >> A;
            tables[i][A]++;
        }
    }

    lld ans = 0;
    rep(i, 0, N) rep(j, i + 1, N)
    {
        ll cnt = 0;
        for (auto &&table : tables[i])
        {
            if (tables[j].contains(table.first))
            {
                cnt += table.second * tables[j][table.first];
            }
        }

        chmax(ans, (lld)cnt / (K[i] * K[j]));
    }

    cout << fix(10) << ans << endl;

    return 0;
}