在R中高效枚举和为常数的多项式组合以计算似然值的技术问询

问题背景 #

给定一个N面非均匀骰子（各面概率为p[1:N]），投掷M次后观测到总和为S，需要计算所有满足和为S的投掷组合的概率乘积之和（即似然值）。直接枚举所有组合（复杂度O(N^M)）在N和M较大时完全不可行，因此需要高效方法。

方法1：动态规划（DP） #

动态规划是最直观的高效解法，核心是逐步构建投掷k次后各可能和的概率总和：

• 状态定义：dp[k][s] 表示投掷k次后和为s的概率总和
• 初始状态：投掷1次时，dp[1][1:N] = p[1:N]，其他和的概率为0
• 递推公式：对于第k次投掷，dp[k][s] = sum(dp[k-1][s - x] * p[x])，其中x需满足 s - x 是k-1次投掷的合法和（即 k-1 ≤ s - x ≤ N*(k-1)）

R代码实现

calc_likelihood_dp <- function(p, M, S) {
  N <- length(p)
  # 边界条件：和不在合法范围内直接返回0
  if (S < M || S > M*N) return(0)
  
  # 初始化DP数组：用向量存储当前次数的各和概率，节省空间
  current_dp <- numeric(M*N)
  current_dp[1:N] <- p  # 1次投掷的情况
  
  for (k in 2:M) {
    next_dp <- numeric(M*N)
    # 遍历k次投掷的所有可能和s
    for (s in k:(k*N)) {
      # 遍历骰子的每个面x，计算s-x是否是k-1次的合法和
      for (x in 1:N) {
        prev_s <- s - x
        if (prev_s >= (k-1) && prev_s <= (k-1)*N) {
          next_dp[s] <- next_dp[s] + current_dp[prev_s] * p[x]
        }
      }
    }
    current_dp <- next_dp
  }
  
  return(current_dp[S])
}

# 测试示例：N=3, M=2, S=4，概率向量假设为c(0.2, 0.5, 0.3)
p_test <- c(0.2, 0.5, 0.3)
calc_likelihood_dp(p_test, 2, 4)
# 预期结果：2*0.2*0.3 + 0.5^2 = 0.12 + 0.25 = 0.37

方法2：生成函数（多项式乘法） #

骰子的概率生成函数为 G(t) = p1*t + p2*t² + ... + pN*t^N，投掷M次的生成函数是 G(t)^M，展开后t^S的系数即为所求似然值。在R中可以用polynom包实现多项式运算，或用convolve函数做卷积加速。

R代码实现（polynom包）

library(polynom)

calc_likelihood_genfunc <- function(p, M, S) {
  N <- length(p)
  if (S < M || S > M*N) return(0)
  
  # 创建概率生成函数多项式：x^1的系数是p[1]，x^2是p[2]，以此类推
  g_poly <- polynomial(coefficients = c(0, p))  # coefficients索引从x^0开始，所以补0在开头
  
  # 计算M次幂
  g_m_poly <- g_poly^M
  
  # 提取x^S的系数（注意polynomial的coefficients是x^0到x^max的系数，所以索引是S+1）
  return(coef(g_m_poly)[S+1])
}

# 测试示例
calc_likelihood_genfunc(p_test, 2, 4)
# 同样返回0.37

效率对比 #

• 动态规划的时间复杂度为O(MS)，当S接近MN时，复杂度接近O(M²N)，但空间可以优化为O(MN)（甚至更小，只保留前一次的结果）
• 生成函数方法依赖多项式乘法，若用快速傅里叶变换（FFT）加速卷积，时间复杂度可降至O(MN log(MN))，适合N和M都较大的场景

两种方法都避免了枚举所有组合，能高效处理较大的N和M值。

内容的提问来源于stack exchange，提问作者Dries