CSES骰子组合问题Python代码提交遇TLE，求优化方案

1. 修复语法错误 #

你的代码中有一个语法错误：\os.fstat(0).st_size 中的反斜杠是多余的，会导致代码无法运行。应改为：

input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

2. 核心优化方向 #

2.1 用滚动和替代重复累加 #

当前代码每次计算dp[i]时都要手动累加6个元素，在n较大时会显著拖慢循环速度。可以维护一个滚动和变量，记录最近6个dp值的总和：

• 每次计算新dp值时直接使用该总和
• 更新总和时，移除最旧的元素并加入新计算的元素，将单次迭代的操作数从6次加法减少到2次算术运算

2.2 减少内存占用 #

不需要存储整个dp数组，只需保留最近6个值即可。使用循环队列（如deque）实现固定大小的缓存，将空间复杂度从O(n)降至O(1)，同时提升数据访问效率。

2.3 函数化代码 #

将逻辑放入main函数中，Python对函数内变量的访问速度更快，能进一步提升运行性能。

优化后的代码示例 #

import sys, io, os
from collections import deque

def main():
    input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
    n = int(input())
    mod = 10**9 +7
    
    if n == 0:
        print(1)
        return
    if n == 1:
        print(1)
        return
    if n <7:
        print((2**(n-1)) % mod)
        return
    
    if n ==6:
        print(32 % mod)
        return
    
    # 初始化队列保存dp[1]到dp[6]的值
    dq = deque([1,2,4,8,16,32])
    current = sum(dq) % mod  # dp[7] = 63
    
    if n ==7:
        print(current)
        return
    
    for _ in range(8, n+1):
        oldest = dq.popleft()
        # 计算下一个值，加mod避免负数
        current = (current *2 - oldest + mod) % mod
        dq.append(current)
    
    print(current % mod)

if __name__ == "__main__":
    main()

关键优化说明 #

• 滚动和计算：通过current *2 - oldest快速得到新的dp值，避免重复累加6个元素，大幅提升循环效率
• 循环队列缓存：用deque的popleft()和append()操作维护最近6个值，保证O(1)的存取速度
• 模运算处理：在减法操作后添加mod确保结果非负，避免出现负数取模的错误

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