如何高效求解满足行长度限制的数组组合最大概率问题

现有一组token，相关参数如下：

• token字符长度列表：

length = [2, 1, 1, 2, 2, 3, 2, 1, 1, 2, 2, 2]

• 每个token对应的[不插入换行符, 插入换行符]概率列表：

prob =  [[9.9978e-01, 2.2339e-04],  [9.9995e-01, 4.9344e-05],  [0.9469, 0.0531],  
         [9.9994e-01, 5.8422e-05],  [0.9964, 0.0036],  [9.9991e-01, 9.4295e-05],  
         [9.9980e-01, 1.9620e-04],  [1.0000e+00, 5.2492e-08],  [9.9998e-01, 1.8293e-05],  
         [9.9999e-01, 5.1220e-06],  [1.0000e+00, 3.9795e-06],  [0.0142, 0.9858]]

• 初始换行方案（仅最后一个token后插入换行）：

[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]

此时总行字符数为21，超出了每行最多20个字符的限制，需要额外插入换行符。最优方案为：

[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1]

该方案选择了插入换行符概率最高的位置，在满足行字符数限制的同时，实现了整体概率乘积的最大化。

当前采用暴力枚举所有2^n种组合（n为token数量）的方式计算概率乘积并筛选，但当n=40时，2^40的计算量导致效率极低，急需高效算法解决。

核心思路 #

我们的目标是找到满足行字符数上限的换行方案，使得所有token的选择概率乘积最大。由于概率乘积的对数等于各概率对数的和，为了避免数值下溢并简化计算，可将乘法问题转换为加法问题——寻找对数概率和最大的合法换行方案。

动态规划状态定义 #

定义dp[i]为处理前i个token时，满足行字符数限制的最大对数概率和；同时维护prev[i]数组，记录达到dp[i]时的前一个换行位置，用于后续回溯生成最终方案。

状态转移与预处理优化 #

前缀和预处理 #

提前计算两个前缀和数组：

1. 长度前缀和：prefix_len[0] = 0，prefix_len[i] = prefix_len[i-1] + length[i-1]，可快速算出从第j+1到第i个token的总长度：prefix_len[i] - prefix_len[j]。
2. 不插入换行的对数概率前缀和：prefix_log_prob0[0] = 0，prefix_log_prob0[i] = prefix_log_prob0[i-1] + math.log(prob[i-1][0])，可快速算出从第j+1到第i-1个token都不插入换行的对数概率和：prefix_log_prob0[i-1] - prefix_log_prob0[j]。

状态转移逻辑 #

遍历每个token位置i（从1到n），向前遍历所有可能的起始位置j（j < i）：

• 若prefix_len[i] - prefix_len[j]不超过行字符上限，则说明从j+1到i的token可以放在同一行，且在i后插入换行。此时计算该方案的对数概率和：

  candidate = dp[j] + (prefix_log_prob0[i-1] - prefix_log_prob0[j]) + math.log(prob[i-1][1])
  

• 若该值大于当前dp[i]，则更新dp[i] = candidate，并记录prev[i] = j。
• 若当前行长度超过上限，即可停止向前遍历j（因为j越小，包含的token越多，总长度只会更大）。

具体实现步骤 #

1. 计算prefix_len和prefix_log_prob0两个前缀和数组。
2. 初始化dp数组：dp[0] = 0，其余元素设为负无穷（表示初始不可达）；prev数组初始化为-1。
3. 遍历i从1到n：
   • 从j = i-1开始向前遍历到0：
     • 计算当前行长度current_len = prefix_len[i] - prefix_len[j]，若current_len > 行上限则break。
     • 计算候选对数概率和candidate，若candidate > dp[i]则更新dp[i]和prev[i]。
4. 回溯prev数组：从i = n开始，依次找到每个换行位置，生成最终的0/1换行方案数组。

复杂度分析 #

每个i最多向前遍历到行长度刚好不超过上限的位置，假设行上限为L，每个token长度至少为1，那么每个i最多遍历L个位置。时间复杂度为O(n*L)，当n=40、L=20时，计算量仅为800，远低于暴力枚举的2^40。

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