作者推荐
本文涉及知识点
深度优先搜索 组合数学
LeetCode1467 两个盒子中球的颜色数相同的概率
桌面上有 2n 个颜色不完全相同的球,球上的颜色共有 k 种。给你一个大小为 k 的整数数组 balls ,其中 balls[i] 是颜色为 i 的球的数量。
所有的球都已经 随机打乱顺序 ,前 n 个球放入第一个盒子,后 n 个球放入另一个盒子(请认真阅读示例 2 的解释部分)。
注意:这两个盒子是不同的。例如,两个球颜色分别为 a 和 b,盒子分别为 [] 和 (),那么 [a] (b) 和 [b] (a) 这两种分配方式是不同的(请认真阅读示例的解释部分)。
请返回「两个盒子中球的颜色数相同」的情况的概率。答案与真实值误差在 10^-5 以内,则被视为正确答案
示例 1:
输入:balls = [1,1]
输出:1.00000
解释:球平均分配的方式只有两种:
- 颜色为 1 的球放入第一个盒子,颜色为 2 的球放入第二个盒子
- 颜色为 2 的球放入第一个盒子,颜色为 1 的球放入第二个盒子
这两种分配,两个盒子中球的颜色数都相同。所以概率为 2/2 = 1 。
示例 2:
输入:balls = [2,1,1]
输出:0.66667
解释:球的列表为 [1, 1, 2, 3]
随机打乱,得到 12 种等概率的不同打乱方案,每种方案概率为 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]
然后,我们将前两个球放入第一个盒子,后两个球放入第二个盒子。
这 12 种可能的随机打乱方式中的 8 种满足「两个盒子中球的颜色数相同」。
概率 = 8/12 = 0.66667
示例 3:
输入:balls = [1,2,1,2]
输出:0.60000
解释:球的列表为 [1, 2, 2, 3, 4, 4]。要想显示所有 180 种随机打乱方案是很难的,但只检查「两个盒子中球的颜色数相同」的 108 种情况是比较容易的。
概率 = 108 / 180 = 0.6 。
提示:
1 <= balls.length <= 8
1 <= balls[i] <= 6
sum(balls) 是偶数
深度优先搜索
极端情况下,8种球,6种颜色。每种球选择0到6个,共7种选择。78 约等于5e6。再加上剪支,能过。
m_iCan 记录,合法选择的可能数。
m_iAns 记录,符合题意的可能数。
注意: 从ball[i]种选择m个求,是组合C b a l l s [ i ] m \Large C_{balls[i]}^mCballs[i]m
代码
核心代码
template<class Result =int > class CCombination { public: CCombination() { m_v.assign(1, vector<Result>(1,1)); } Result Get(int sel, int total) { while (m_v.size() <= total) { int iSize = m_v.size(); m_v.emplace_back(iSize + 1, 1); for (int i = 1; i < iSize; i++) { m_v[iSize][i] = m_v[iSize - 1][i] + m_v[iSize - 1][i - 1]; } } return m_v[total][sel]; } protected: vector<vector<Result>> m_v; }; class Solution { public: double getProbability(vector<int>& balls) { m_iN = std::accumulate(balls.begin(), balls.end(), 0) / 2; DFS(balls, 0, 0, 0, 0,1); return (double)m_iiAns / m_iiSel; } void DFS(const vector<int>& balls,int iCur,int iHasSel,int iSelAll,int iSel0,long long iiMul) { if (iHasSel == m_iN) { m_iiSel += iiMul; if (iSelAll == iSel0 + balls.size()- iCur ) {//余下的球全部不选择 m_iiAns += iiMul; } return; } if (iCur >= balls.size()) { return ; } for (int curSel = 0; (curSel <= balls[iCur])&&(curSel+iHasSel <= m_iN); curSel++) { DFS(balls, iCur + 1, curSel + iHasSel, iSelAll + (curSel == balls[iCur]), iSel0 + (0 == curSel),iiMul*m_com.Get(curSel, balls[iCur])); } } long long m_iN, m_iiSel=0, m_iiAns=0; CCombination<int> m_com; };
测试用例
template<class T> void Assert(const T& t1, const T& t2) { assert(t1 == t2); } template<class T> void Assert(const vector<T>& v1, const vector<T>& v2) { if (v1.size() != v2.size()) { assert(false); return; } for (int i = 0; i < v1.size(); i++) { Assert(v1[i], v2[i]); } } int main() { vector<int> balls; { Solution sln; balls = { 1, 1 }; auto res = sln.getProbability(balls); assert(abs(res - 1 ) < 0.0001); } { Solution sln; balls = { 2,1,1 }; auto res = sln.getProbability(balls); assert(abs(res - 0.66667) < 0.0001); } { Solution sln; balls = { 1,2,1,2 }; auto res = sln.getProbability(balls); assert(abs(res - 0.6) < 0.0001); } { Solution sln; balls = { 6, 6, 6, 6, 6, 6, 6, 6 }; auto res = sln.getProbability(balls); assert(abs(res - 0.85571) < 0.0001); } }
动态规划
动态规划的状态表示
pre[sel][c]记录可能排列数量。sel表示第一个盒子的球数,c表示颜色差。c等于0,表示左边全选的球的数量 比 右边全先的求的数量 少6。 c = 全部在第一个盒子的颜色数- 全部在第二个盒子的颜色+6。
不在两种颜色相差8的情况:那样一个盒子为空,和n个球矛盾。
不存在颜色相差7的情况:全选7种颜色,至少有7个球。全先1种颜色顶多6个球。无法相等。
存在相差6的情况:{** 1 1 1 1 1 1 ** 3 3} 。前6个球是1,全选。
class Solution { public: double getProbability(vector<int>& balls) { const int n = std::accumulate(balls.begin(), balls.end(), 0) / 2; vector<vector<long long>> pre(n + 1, vector<long long>(13, 0)); pre[0][6] = 1; for (const auto& b : balls) { vector<vector<long long>> dp(n + 1, vector<long long>(13, 0)); for (int col = 0; col < 13; col++) { for (int preSel = 0; preSel <= n; preSel++) { for (int curSel = 0; (curSel <= b) && (preSel + curSel <= n); curSel++) { int col1 = col + (curSel == b) - (curSel == 0); if ((col1 >= 0) && (col1 < 13)) { dp[preSel + curSel][col1] += pre[preSel][col]*m_com.Get(curSel,b); } } } } pre.swap(dp); } long long llAns = pre.back()[6], llSel = std::accumulate(pre.back().begin(), pre.back().end(),0LL); return (double)llAns / llSel; } CCombination<int> m_com; };
2023年2月版
class Solution {
public:
double getProbability(const vector& balls) {
const int iTotal = std::accumulate(balls.begin(), balls.end(), 0);
m_c = balls.size();
vector<vector> combinations(6 + 1, vector(6 + 1, 1));
for (int i = 1; i <= 6; i++)
{
for (int j = 1; j < i; j++)
{
combinations[i][j] = combinations[i - 1][j - 1] + combinations[i - 1][j];
}
}
vector<vector> pre(13, vector(iTotal + 1));
pre[6][0] = 1;
for (int i = 0; i < balls.size(); i++)
{
vector<vector> dp(13, vector(iTotal + 1));
for (int colorDiff = 0; colorDiff < 13; colorDiff++)
{
for (int selBallNum = 0; selBallNum <= iTotal; selBallNum++)
{
if (0 == pre[colorDiff][selBallNum])
{
continue;
}
for (int k = 0; k <= balls[i]; k++)
{
int iNewColorDiff = colorDiff;
if (0 == k)
{
iNewColorDiff–;
}
if (balls[i] == k)
{
iNewColorDiff++;
}
if ((iNewColorDiff<0) || (iNewColorDiff >12))
{
continue;
}
const int iNewSelBallNum = selBallNum + k;
if ( iNewSelBallNum > iTotal)
{
continue;
}
dp[iNewColorDiff][iNewSelBallNum] += pre[colorDiff][selBallNum] * combinations[balls[i]][k];
}
}
}
pre.swap(dp);
}
double dNum = 0, dEqualNum = 0;
for (int colorDiff = 0; colorDiff < 13; colorDiff++)
{
const int selBallNum = iTotal / 2;
//for (int selBallNum = 0; selBallNum <= iTotal; selBallNum++)
{
const double dAdd = (double)pre[colorDiff][selBallNum] ;
dNum += dAdd;
if (6 == colorDiff)
{
dEqualNum += dAdd;
}
}
}
return (double)dEqualNum / dNum;
}
int m_c;
};