基于记忆化优化网格皇后路径计数的C++代码优化求助

你的原代码核心问题是重复累加计算：每个单元格dp[i][j]的值需要遍历左侧、上方、对角线的所有前置单元格求和，时间复杂度达到O(n*m*(n+m))，当n或m较大时性能会急剧下降。我们可以通过维护前缀和数组来替代逐一遍历求和，将时间复杂度优化到O(n*m)，同时保留动态规划（记忆化）的核心逻辑。

优化思路 #

我们需要三个辅助前缀和数组，分别记录：

• 行前缀和：row_sum[i] 表示第i行到当前列之前的所有dp值之和，用于快速计算从左侧任意位置跳到(i,j)的路径总数。
• 列前缀和：col_sum[j] 表示第j列到当前行之前的所有dp值之和，用于快速计算从上方任意位置跳到(i,j)的路径总数。
• 对角线前缀和：diag_sum[k] 表示对应对角线（满足i-j = 固定值）到当前位置之前的所有dp值之和，用于快速计算从左上对角线任意位置跳到(i,j)的路径总数。

通过这三个数组，我们可以在O(1)时间内得到每个方向的累加和，避免重复遍历。

优化后的完整代码 #

#include <iostream>
#include <vector>

long long count_ways(int n, int m) {
    // dp[i][j] 表示从(0,0)到(i,j)的路径数
    std::vector<std::vector<long long>> dp(n, std::vector<long long>(m, 0));
    dp[0][0] = 1;

    // 行前缀和：row_sum[i] 记录第i行前j个元素的和（更新dp[i][j]前的状态）
    std::vector<long long> row_sum(n, 0);
    row_sum[0] = dp[0][0];

    // 列前缀和：col_sum[j] 记录第j列前i个元素的和（更新dp[i][j]前的状态）
    std::vector<long long> col_sum(m, 0);
    col_sum[0] = dp[0][0];

    // 对角线前缀和：对角线满足i-j = d，d范围是-(m-1)到n-1，偏移m-1转为非负索引
    std::vector<long long> diag_sum(n + m - 1, 0);
    diag_sum[0 + m - 1] = dp[0][0];

    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < m; ++j) {
            // 跳过起点，已经初始化
            if (i == 0 && j == 0) continue;

            long long total = 0;
            // 左侧路径和：取当前行的前缀和
            if (j > 0) {
                total += row_sum[i];
            }
            // 上方路径和：取当前列的前缀和
            if (i > 0) {
                total += col_sum[j];
            }
            // 对角线路径和：取对应对角线的前缀和
            if (i > 0 && j > 0) {
                int diag_idx = (i - j) + m - 1;
                total += diag_sum[diag_idx];
            }

            dp[i][j] = total;

            // 更新前缀和数组
            row_sum[i] += dp[i][j];
            col_sum[j] += dp[i][j];
            int diag_idx = (i - j) + m - 1;
            diag_sum[diag_idx] += dp[i][j];
        }
    }

    return dp[n - 1][m - 1];
}

int main() {
    int n = 3, m = 3;
    std::cout << count_ways(n, m) << std::endl; // 输出22，符合预期
    return 0;
}

关键细节说明 #

1. 前缀和的更新时机：每次计算完dp[i][j]后，立即更新对应的行、列、对角线前缀和，确保后续单元格能获取到最新的累加值。
2. 对角线索引处理：由于i-j可能为负数，我们通过偏移m-1将索引转为非负值，避免数组越界。
3. 避免全局变量：将dp和前缀和数组放在函数内部，避免多次调用时的状态残留问题。
4. 数据类型：使用long long防止路径数过大导致整数溢出。

性能对比 #

• 原代码：时间复杂度O(n*m*(n+m))，空间复杂度O(n*m)
• 优化后代码：时间复杂度O(n*m)，空间复杂度O(n*m + n + m)（额外的前缀和数组开销可以忽略）

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