Skip to content

Probability DP

Introduction

Probability DP is used to solve probability and expectation problems. It is recommended to have some understanding of Probability & Expectation first. Generally, solving probability problems requires forward iteration, while solving expectation problems uses backward iteration. If the defined state transition equation has aftereffects, Gaussian elimination is needed for optimization. Probability DP may also be combined with other knowledge, such as state compression and DP transitions on trees.

Probability DP

These types of problems use forward iteration, meaning they transition from the initial state to the result. Similar to general DP, the difficulty still lies in formulating the state transition equation, but these problems are wrapped in probability theory knowledge.

Example

Codeforces 148D Bag of mice

There are \(w\) white mice and \(b\) black mice in a bag. The princess and the dragon take turns drawing mice from the bag. Whoever draws a white mouse first wins. If there are no mice left in the bag and neither has drawn a white mouse, the dragon wins. The princess draws one mouse each time, and after the dragon draws a mouse, one mouse runs out of the bag. Each drawn mouse and the mouse that runs out are random. The princess draws first. What is the probability that the princess wins?

Solution

Let \(f_{i,j}\) be the probability that the princess wins when it is her turn and there are \(i\) white mice and \(j\) black mice in the bag. Initialize the boundaries: \(f_{0,j}=0\) because if there are no white mice, the dragon wins; \(f_{i,0}=1\) because drawing one mouse will be white, so the princess wins. Consider the transition for \(f_{i,j}\):

  • The princess draws a white mouse and wins. Probability: \(\dfrac{i}{i+j}\).
  • The princess draws a black mouse, the dragon draws a white mouse, and the dragon wins. Probability: \(\dfrac{j}{i+j}\cdot\dfrac{i}{i+j-1}\).
  • The princess draws a black mouse, the dragon draws a black mouse, and a black mouse runs out, transitioning to \(f_{i,j-3}\). Probability: \(\dfrac{j}{i+j}\cdot\dfrac{j-1}{i+j-1}\cdot\dfrac{j-2}{i+j-2}\).
  • The princess draws a black mouse, the dragon draws a black mouse, and a white mouse runs out, transitioning to \(f_{i-1,j-2}\). Probability: \(\dfrac{j}{i+j}\cdot\dfrac{j-1}{i+j-1}\cdot\dfrac{i}{i+j-2}\).

Consider the probability of the princess winning; the second case does not contribute. Also, ensure the latter two cases are valid: for the third case, at least 3 black mice are needed; for the fourth case, at least 1 white mouse and 2 black mice are needed.

Reference Code 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#include <cstring>
#include <iomanip>
#include <iostream>
using namespace std;

using ll = long long;
int w, b;
double dp[1010][1010];

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> w >> b;
  memset(dp, 0, sizeof(dp));
  for (int i = 1; i <= w; i++) dp[i][0] = 1;  // Initialization
  for (int i = 1; i <= b; i++) dp[0][i] = 0;
  for (int i = 1; i <= w; i++) {
    for (int j = 1; j <= b; j++) {  // Probability transition according to the problem statement
      dp[i][j] += (double)i / (i + j);
      if (j >= 3) {
        dp[i][j] += (double)j / (i + j) * (j - 1) / (i + j - 1) * (j - 2) /
                    (i + j - 2) * dp[i][j - 3];
      }
      if (i >= 1 && j >= 2) {
        dp[i][j] += (double)j / (i + j) * (j - 1) / (i + j - 1) * i /
                    (i + j - 2) * dp[i - 1][j - 2];
      }
    }
  }
  cout << fixed << setprecision(9) << dp[w][b] << '\n';
  return 0;
}

Exercises

Expectation DP

Example

POJ2096 Collecting Bugs

A software has \(s\) subsystems and produces \(n\) types of bugs. Someone finds one bug per day. Each bug belongs to a bug category and a subsystem. The probability that a bug belongs to a particular subsystem is \(\dfrac{1}{s}\), and the probability that it belongs to a particular bug category is \(\dfrac{1}{n}\). Find the expected number of days to discover all \(n\) bug categories and bugs in all \(s\) subsystems.

Solution

Let \(f_{i,j}\) be the expected number of days to reach the target state after having found \(i\) bug categories and bugs in \(j\) subsystems. The target state is to find all \(n\) bug categories and bugs in all \(s\) subsystems. Then \(f_{n,s}=0\), as the target state is already reached and no more days are needed to discover bugs. Thus, we start the recurrence from the target state, and the answer is \(f_{0,0}\).

Consider the state transition for \(f_{i,j}\):

  • \(f_{i,j}\): discover a bug belonging to the already found \(i\) bug categories and \(j\) subsystems. Probability: \(p_1=\dfrac{i}{n}\cdot\dfrac{j}{s}\).
  • \(f_{i,j+1}\): discover a bug belonging to the already found \(i\) bug categories but not to the already found subsystems. Probability: \(p_2=\dfrac{i}{n}\cdot(1-\dfrac{j}{s})\).
  • \(f_{i+1,j}\): discover a bug not belonging to the already found bug categories but belonging to \(j\) subsystems. Probability: \(p_3=(1-\dfrac{i}{n})\cdot\dfrac{j}{s}\).
  • \(f_{i+1,j+1}\): discover a bug not belonging to the already found bug categories and not belonging to the already found subsystems. Probability: \(p_4=(1-\dfrac{i}{n})\cdot(1-\dfrac{j}{s})\).

Using the linearity of expectation, we obtain the state transition equation:

\[ \begin{aligned} f_{i,j} &= p_1\cdot f_{i,j}+p_2\cdot f_{i,j+1}+p_3\cdot f_{i+1,j}+p_4\cdot f_{i+1,j+1} + 1\\ &= \dfrac{p_2\cdot f_{i,j+1}+p_3\cdot f_{i+1,j}+p_4\cdot f_{i+1,j+1}+1}{1-p_1} \end{aligned} \]
Reference Code 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
#include <iomanip>
#include <iostream>
using namespace std;
int n, s;
double dp[1010][1010];

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n >> s;
  dp[n][s] = 0;
  for (int i = n; i >= 0; i--) {
    for (int j = s; j >= 0; j--) {
      if (i == n && s == j) continue;
      dp[i][j] = (dp[i][j + 1] * i * (s - j) + dp[i + 1][j] * (n - i) * j +
                  dp[i + 1][j + 1] * (n - i) * (s - j) + n * s) /
                 (n * s - i * j);  // Probability transition
    }
  }
  cout << fixed << setprecision(4) << dp[0][0] << '\n';
  return 0;
}
「NOIP2016」Classroom Change

Niu Niu has \(n\) classes in \(n\) time slots. The \(i\)-th class is in classroom \(c_i\), but he can apply to change to classroom \(d_i\) with success probability \(p_i\). He can apply for at most \(m\) class changes. After the \(i\)-th class, he must walk to the classroom for the \((i+1)\)-th class. Given a graph with \(v\) classrooms and \(e\) roads, moving consumes physical energy. Determine which classes to apply for to minimize the expected total physical energy consumed from moving between classrooms, i.e., find the minimum expected total distance.

Solution

For this undirected connected graph, first use Floyd's algorithm to find the shortest paths, which will facilitate subsequent state transitions. Treat moving one step as a stage (moving from the \(i\)-th time slot to the \((i+1)\)-th time slot is one step). In each step, there is a probability \(p_i\) of going to \(d_i\) (but only \(m\) of the \(d_i\) can be chosen overall) and a probability \(1-p_i\) of going to \(c_i\). Find the minimum expected total distance after \(n\) stages.

Define \(f_{i,j,0/1}\) as the minimum expected total distance up to the \(i\)-th time slot, having used \(j\) classroom change opportunities so far (including this slot), and whether a change is made in this slot (1) or not (0). The answer is \(\min \{f_{n,i,0},f_{n,i,1}\} ,i\in[0,m]\). Note the boundaries: \(f_{1,0,0}=f_{1,1,1}=0\).

Consider the state transition for \(f_{i,j,0/1}\):

  • If no change is made in this stage, i.e., \(f_{i,j,0}\). It could come from the previous state with no change: \(f_{i-1,j,0}+w_{c_{i-1},c_{i}}\). Or it could come from the previous state with a change. Combining conditional probability and total probability, we get: \(f_{i-1,j,1}+w_{d_{i-1},c_{i}}\cdot p_{i-1}+w_{c_{i-1},c_{i}}\cdot (1-p_{i-1})\). Thus, the state transition equation is:
\[ \begin{aligned} f_{i,j,0}=min(f_{i-1,j,0}+w_{c_{i-1},c_{i}},f_{i-1,j,1}+w_{d_{i-1},c_{i}}\cdot p_{i-1}+w_{c_{i-1},c_{i}}\cdot (1-p_{i-1})) \end{aligned} \]
  • If a change is made in this stage, i.e., \(f_{i,j,1}\). Similarly, it could come from the previous state with no change or from the previous state with a change. For no change, multiply by \((1-p_i)\); for a change, multiply by \(p_i\). Enumerate all possible cases and compute accordingly. The various transition cases are not elaborated here, but with the previous example, the state transition should be easy to write.
Reference Code 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#include <algorithm>
#include <iomanip>
#include <iostream>

using namespace std;

constexpr int MAXN = 2010;
int n, m, v, e;
int f[MAXN][MAXN], c[MAXN], d[MAXN];
double dp[MAXN][MAXN][2], p[MAXN];

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> n >> m >> v >> e;
  for (int i = 1; i <= n; i++) cin >> c[i];
  for (int i = 1; i <= n; i++) cin >> d[i];
  for (int i = 1; i <= n; i++) cin >> p[i];
  for (int i = 1; i <= v; i++)
    for (int j = 1; j < i; j++) f[i][j] = f[j][i] = 1e9;

  int u, V, w;
  for (int i = 1; i <= e; i++) {
    cin >> u >> V >> w;
    f[u][V] = f[V][u] = min(w, f[u][V]);
  }

  for (int k = 1; k <= v; k++)
    for (int i = 1; i <= v; i++)  // Previous, perform state transition according to the previous solution
      for (int j = 1; j < i; j++)
        if (f[i][k] + f[k][j] < f[i][j]) f[i][j] = f[j][i] = f[i][k] + f[k][j];

  for (int i = 1; i <= n; i++)
    for (int j = 0; j <= m; j++) dp[i][j][0] = dp[i][j][1] = 1e9;

  dp[1][0][0] = dp[1][1][1] = 0;
  for (int i = 2; i <= n; i++)  // Equation with aftereffects
    for (int j = 0; j <= min(i, m); j++) {
      dp[i][j][0] = min(dp[i - 1][j][0] + f[c[i - 1]][c[i]],
                        dp[i - 1][j][1] + f[c[i - 1]][c[i]] * (1 - p[i - 1]) +
                            f[d[i - 1]][c[i]] * p[i - 1]);
      if (j != 0) {
        dp[i][j][1] = min(dp[i - 1][j - 1][0] + f[c[i - 1]][d[i]] * p[i] +
                              f[c[i - 1]][c[i]] * (1 - p[i]),
                          dp[i - 1][j - 1][1] +
                              f[c[i - 1]][c[i]] * (1 - p[i - 1]) * (1 - p[i]) +
                              f[c[i - 1]][d[i]] * (1 - p[i - 1]) * p[i] +
                              f[d[i - 1]][c[i]] * (1 - p[i]) * p[i - 1] +
                              f[d[i - 1]][d[i]] * p[i - 1] * p[i]);
      }
    }

  double ans = 1e9;
  for (int i = 0; i <= m; i++) ans = min(dp[n][i][0], min(dp[n][i][1], ans));
  cout << fixed << setprecision(2) << ans;

  return 0;
}

Comparing these two problems, we see that DP for expectation problems may affect the transition method depending on whether a specific value is sought or an optimization is performed. However, whether DP is used for probability or expectation, the steps always involve probability knowledge and formulating/simplifying calculation formulas, and similar details need to be considered when writing state transition equations.

Exercises

DP with Aftereffects

Example

CodeForces 24D Broken robot

Given an \(n \times m\) grid. A robot starts at row \(x\) and column \(y\). Each step, the robot chooses with equal probability to stay in place, move left, move right, or move down. If the robot is at the boundary, it will not move outside the grid. Find the expected number of steps for the robot to reach the last row.

Solution

When \(m=1\), each step has a \(\dfrac{1}{2}\) probability of staying and a \(\dfrac{1}{2}\) probability of moving down one row. The answer is \(2\cdot (n-x)\). Let \(f_{i,j}\) be the expected number of steps for the robot to reach row \(n\) starting from row \(i\) and column \(j\). The final state is \(f_{n,j}=0\). Since the robot chooses with equal probability to stay, move left, move right, or move down, consider the state transition for \(f_{i,j}\):

  • \(f_{i,1}=\dfrac{1}{3}\cdot(f_{i+1,1}+f_{i,2}+f_{i,1})+1\)
  • \(f_{i,j}=\dfrac{1}{4}\cdot(f_{i,j}+f_{i,j-1}+f_{i,j+1}+f_{i+1,j})+1\)
  • \(f_{i,m}=\dfrac{1}{3}\cdot(f_{i,m}+f_{i,m-1}+f_{i+1,m})+1\)

Between rows, since the robot can only move down, there are no aftereffects. Between columns, the robot can move left or right, which may create cycles during movement, leading to aftereffects. Transforming the equations gives:

  • \(2f_{i,1}-f_{i,2}=3+f_{i+1,1}\)
  • \(3f_{i,j}-f_{i,j-1}-f_{i,j+1}=4+f_{i+1,j}\)
  • \(2f_{i,m}-f_{i,m-1}=3+f_{i+1,m}\)

Since the recurrence is backward, each \(f_{i+1,j}\) is known. With \(m\) columns, the right side is essentially a column vector of \(m\) rows, and the left side is an \(m \times m\) matrix. Using an augmented matrix, it becomes an \(m \times (m+1)\) matrix, and then Gaussian elimination can be applied to solve for the answer.

Reference Code 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
#include <cstdio>
#include <cstring>
using namespace std;

constexpr int MAXN = 1e3 + 10;

double a[MAXN][MAXN], f[MAXN];
int n, m;

void solve(int x) {
  for (int i = 1; i <= m; i++) {
    if (i == 1) {
      a[i][i] = 2;
      a[i][i + 1] = -1;
      a[i][m + 1] = 3 + f[i];
      continue;
    } else if (i == m) {
      a[i][i] = 2;
      a[i][i - 1] = -1;
      a[i][m + 1] = 3 + f[i];
      continue;
    }
    a[i][i] = 3;
    a[i][i + 1] = -1;
    a[i][i - 1] = -1;
    a[i][m + 1] = 4 + f[i];
  }

  for (int i = 1; i < m; i++) {
    double p = a[i + 1][i] / a[i][i];
    a[i + 1][i] = 0;
    a[i + 1][i + 1] -= a[i][i + 1] * p;
    a[i + 1][m + 1] -= a[i][m + 1] * p;
  }

  f[m] = a[m][m + 1] / a[m][m];
  for (int i = m - 1; i >= 1; i--)
    f[i] = (a[i][m + 1] - f[i + 1] * a[i][i + 1]) / a[i][i];
}

int main() {
  scanf("%d %d", &n, &m);
  int st, ed;
  scanf("%d %d", &st, &ed);
  if (m == 1) {
    printf("%.10f\n", 2.0 * (n - st));
    return 0;
  }
  for (int i = n - 1; i >= st; i--) {
    solve(i);
  }
  printf("%.10f\n", f[ed]);
  return 0;
}

Exercises

References

kuangbin's Probability DP Summary

本页面最近更新：，更新历史
发现错误？想一起完善？ 在 GitHub 上编辑此页！
本页面贡献者：OI-wiki
本页面的全部内容在 CC BY-SA 4.0 和 SATA 协议之条款下提供，附加条款亦可能应用