LeetCode Weekly 191 - D (組み合わせ数)

組み合わせ数（確率）を求める問題。苦戦したので解法メモを残しておく。
https://leetcode.com/problems/probability-of-a-two-boxes-having-the-same-number-of-distinct-balls/

色1~kで染められたボールが2n個ある。ランダムにn個ずつ2グループに分けたとき、それぞれのグループにある色の種類数が同じになる確率を求めよ。

1 <= k <= 8
1 <= 各色のボールの数 <= 6

たとえば次の場合、
色1： 2個
色2： 1個
色3： 1個
組み合わせは12通りになる。
[1,1 / 2,3], [1,1 / 3,2], [1,2 / 1,3], [1,2 / 3,1], [1,3 / 1,2], [1,3 / 2,1], [2,1 / 1,3], [2,1 / 3,1], [2,3 / 1,1], [3,1 / 1,2], [3,1 / 2,1], [3,2 / 1,1]
このうち8パターンで、各グループにある色の種類数が同じになる。
[1,2 / 1,3], [1,2 / 3,1], [1,3 / 1,2], [1,3 / 2,1], [2,1 / 1,3], [2,1 / 3,1], [3,1 / 1,2], [3,1 / 2,1]
よって確率は8/12=0.6667。

考え方
順番は無視できる。上の例だと
[1,1 / 2,3], [1,2 / 1,3], [1,3 / 1,2], [2,1 / 1,3], [2,3 / 1,1], [3,1 / 1,2],
[1,2 / 1,3], [1,3 / 1,2], [2,1 / 1,3], [3,1 / 1,2],
で4/6=0.6667。

まず分母は全組み合わせ数なのでnCkで求められる。

次にグループ名をA・Bとしたとき、dp[i, j, k, l]を
i: Aに入っているボールの数
j: Bに入っているボールの数
k: 箱Aの色種数
l: 箱Bの色種数
の組み合わせ数と定義する。

色cのボールがn個あるとすると、x: 0～nについて
dp[i + x, j + n - x, k + 1, l + 1] += dp[i, j, k, l] * nCk(n, x);
と更新していくことができる。

最後にdpテーブルから条件に合うものを合算し、分子が求まる。

コード例は配列ではなくDictionaryを使った。

public class Solution
{
    public double GetProbability(int[] balls)
    {
        var nCk = GetNcks();
        var sum = balls.Sum();

        // 分母。全ボールから半分選ぶパターン数
        var denominator = GetNck(sum, sum / 2);

        // i: 箱Aに入っているボールの数
        // j: 箱Bに入っているボールの数
        // k: 箱Aの色種数
        // l: 箱Bの色種数
        var dp = new Dictionary<(int i, int j, int k, int l), long>();
        dp[(0, 0, 0, 0)] = 1;

        // c: 色の種類
        for (int c = 0; c < balls.Length; c++)
        {
            var newDp = new Dictionary<(int i, int j, int k, int l), long>();

            var n = balls[c];
            for (int x = 0; x <= n; x++)
            {
                // ボールn個からx個、箱Aへ入れる

                foreach(var key in dp.Keys)
                {
                    if (x == 0)
                    {
                        if (!newDp.ContainsKey((key.i, key.j + n, key.k, key.l + 1))) newDp[(key.i, key.j + n, key.k, key.l + 1)] = 0;
                        newDp[(key.i, key.j + n, key.k, key.l + 1)] += dp[key] * nCk[(n, 0)];
                    }
                    else if (x == n)
                    {
                        if (!newDp.ContainsKey((key.i + n, key.j, key.k + 1, key.l))) newDp[(key.i + n, key.j, key.k + 1, key.l)] = 0;
                        newDp[(key.i + n, key.j, key.k + 1, key.l)] += dp[key] * nCk[(n, n)];
                    }
                    else
                    {
                        if (!newDp.ContainsKey((key.i + x, key.j + n - x, key.k + 1, key.l + 1))) newDp[(key.i + x, key.j + n - x, key.k + 1, key.l + 1)] = 0;

                        // 次の状態：箱Aにi+x個、箱Bにj+n-x個、箱Aにk+1種類、箱Bにl+1種類
                        // nからx選ぶパターン数nCxを掛けて足す
                        newDp[(key.i + x, key.j + n - x, key.k + 1, key.l + 1)] += dp[key] * nCk[(n, x)];
                    }
                }
            }
            dp = newDp;
        }

        // 分子。条件を満たす組み合わせを合算
        var numerator = 0L;
        foreach(var key in dp.Keys)
        {
            if (key.Item1 == key.Item2 && key.Item3 == key.Item4)
                numerator += dp[key];
        }

        return (double)numerator / denominator;
    }

    static Dictionary<(int, int), long> GetNcks()
    {
        var ret = new Dictionary<(int, int), long>();

        // max 6 balls in each color
        for(int i=1; i<=6; i++)
        {
            for(int j=0; j<=i; j++)
            {
                ret[(i, j)] = GetNck(i, j);
            }
        }
        return ret;
    }

    public static long GetNck(int n, int k)
    {
        long val;
        if (n < k) return 0;
        k = Math.Min(k, n - k);
        if (k == 0) val = 1;
        else val = GetNck(n - 1, k - 1) * n / k;
        return val;
    }
}