绘画栅栏问题代码大输入输出错误：n=27608,k=99118结果不符

你编写的绘画栅栏问题C语言代码在多数场景下输出正确，但输入n=27608、k=99118时结果错误，核心原因是初始化b[0]时发生了整数溢出。

问题代码 #

long long countWays(int n, int k){
    // code here
    if(n==1) return k;
    long long a[n-1];
    long long b[n-1];
    // for n=2
    a[0]=k;
    b[0]=k*(k-1); // 此处存在溢出风险
    for(int i=1;i<n-1;i++){
        a[i] = b[i-1]%1000000007;
        b[i] = ((a[i-1] + b[i-1])*(k-1))%1000000007;
    }
    return (a[n-2] + b[n-2]) %1000000007;
}

错误原因分析 #

k是int类型，当k=99118时，k*(k-1)的计算结果为99118*99117=9824377926，这个数值远超过了int类型的最大取值范围（通常为2^31-1=2147483647）。此时乘法运算会先以int类型执行，直接触发整数溢出，得到错误的截断值，之后再赋值给long long变量也无法修正这个错误结果。

修复方案 #

将k强制转换为long long后再执行乘法，确保计算过程在64位整数范围内进行：

b[0] = (long long)k * (k-1);

另外，代码中用数组存储所有中间值会占用不必要的内存（当n=27608时，两个数组共占用约220KB内存），可以优化为仅保存前一轮的a和b值，进一步节省空间：

long long countWays(int n, int k){
    if(n==1) return k;
    long long prev_a = k;
    long long prev_b = (long long)k * (k-1);
    for(int i=2;i<n;i++){
        long long curr_a = prev_b % 1000000007;
        long long curr_b = ((prev_a + prev_b) * (k-1)) % 1000000007;
        prev_a = curr_a;
        prev_b = curr_b;
    }
    return (prev_a + prev_b) % 1000000007;
}

内容的提问来源于stack exchange，提问作者Swadhin Ranjan Patra