拼多多笔试题：多多君住宿舍（概率DP）

思路：概率dp，我们设dp[i][j][k]：表示前i天多多君有j天使一个人过夜并且第i天宿舍有k个人的概率。

public class Solution {

    static int n, m;
    static double a, b;

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        n = in.nextInt();
        m = in.nextInt();
        a = in.nextDouble();
        b = in.nextDouble();
        if (n == 1 || m == 1) {
            double x = 1.0;
            System.out.printf("%.8f
", x);
        } else {
            double[][][] dp = new double[m + 1][m + 1][n + 1];
            dp[1][1][1] = 1.0;
            for (int d = 2; d <= m; d++) {
                for (int i = 1; i <= d; i++) {
                    for (int j = 1; j <= n; j++) {
                        for (int k = 1; k <= n; k++) {
                            if (j == 1)
                                dp[d][i][j] += dp[d - 1][i - 1][k] * work1(j, k);
                            else
                                dp[d][i][j] += dp[d - 1][i][k] * work2(j, k);
                        }
                    }
                }
            }
            double ans = 0;
            for (int i = 1; i <= m; i++)
                for (int j = 1; j <= n; j++)
                    ans += dp[m][i][j] * i;
            System.out.printf("%.8f
", ans);
        }
    }

    private static double work1(int x, int y) {
        if (x == y)
            return 1.0 - a;
        else {
            if (y - x == 1)
                return n == 2 ? b : (1.0 - a) * b;
            else
                return b * b;
        }
    }

    private static double work2(int x, int y) {
        if (x == 2 && y == 1) return a;
        if (x == 2 && y == 2) return n == 2 ? (1.0 - b) : (1.0 - b) * (1.0 - a) + a * b;
        if (x == 2 && y == 3) return 2.0 * (1.0 - b) * b;
        if (x == 3 && y == 1) return 0;
        if (x == 3 && y == 2) return a * (1.0 - b);
        return (1.0 - b) * (1.0 - b);
    }
}