# 2. 构造兼容性矩阵:若两列同一行颜色都不相等,便兼容 compatible = [[True] * S for _ inrange(S)] for i inrange(S): for j inrange(S): # 检查 states[i] 和 states[j] 是否在每一行都不同色 for row inrange(m): if states[i][row] == states[j][row]: compatible[i][j] = False break
# 3. 动态规划(滚动数组)仅保留上一列的 dp # dp_prev[i] = 以第 c-1 列状态 i 的方案数;dp_cur[i] = 以第 c 列状态 i 的方案数 dp_prev = [1] * S # 第一列每个状态方案数 = 1 dp_cur = [0] * S
for _ inrange(2, n + 1): # 每次计算第 col 列 for j inrange(S): total = 0 for i inrange(S): if compatible[i][j]: total += dp_prev[i] dp_cur[j] = total % self.MOD # 滚动:更新上一列数组 dp_prev, dp_cur = dp_cur, [0] * S
defdfs_build(col: int, prev_color: int): # col : 当前要填到第几行(0-based) # prev_color: 上一行填的是什么颜色(若为 -1,表示还没填,即当前处于第一行) if col == m: # path 长度正好等于 m,且满足相邻不同色,记录一个合法状态 states.append(tuple(path)) return # 枚举三种颜色 for c in (0, 1, 2): if c != prev_color: path.append(c) dfs_build(col + 1, c) path.pop()
dfs_build(0, -1) # 合法状态总数 S = len(states) # states[i] 就是第 i 个合法状态,对应一个长度 m 的 0/1/2 元组 # 这个 S 一定等于 3 * 2^(m-1)。当 m=5 时,S=48;m=1 时,S=3;以此类推。
# ----------------------------------------------------------- # 2. 构造“邻接表”:对每一个状态 i,预先把所有与之“列间颜色都不同”的 j 放到 adj[i] 里 # 这样 DP 转移时,直接遍历 adj[i],而不用再在循环里做兼容性判断。 # 兼容条件:对所有 0 <= row < m,states[i][row] != states[j][row] # ----------------------------------------------------------- adj = [[] for _ inrange(S)] for i inrange(S): si = states[i] for j inrange(S): sj = states[j] # 检查 si 和 sj 对应行是否都不相同 ok = True for row inrange(m): if si[row] == sj[row]: ok = False break if ok: adj[i].append(j)
# ----------------------------------------------------------- # 3. 动态规划(滚动数组) # dp_prev[i] 表示:上一列(第 c-1 列)选状态 i 时的方案数 # dp_cur[i] 表示:当前列 (第 c 列)选状态 i 时的方案数 # # 初始条件 (c=1):第 1 列任何状态都可以,记为 1。 # 转移 (c -> c+1):dp_cur[j] = sum{ dp_prev[i] | i in adj[j] } % MOD # 最后答案 = sum(dp_prev[i] for i in 0..S-1) % MOD # # 关键优化点在于:我们 **只遍历 adj[j]** 而不是所有 i,再去判断兼容与否, # 从 O(S^2 * m) 的常数提到 O(Σ|adj[j]|),非常节省时间。 # ----------------------------------------------------------- dp_prev = [1] * S # c=1 的时候,每个状态 i 都只有一种选择 dp_cur = [0] * S
# 从第 2 列开始,依次做转移 for _ inrange(2, n + 1): # 遍历“本列”可能的状态 j for j inrange(S): s = 0 # 只遍历可以跟 j 配对的 i 列表 for i in adj[j]: s += dp_prev[i] # 取模后写入 dp_cur[j] = s % MOD
# 滚动:把当前列当作下一轮的“上一列” dp_prev, dp_cur = dp_cur, [0] * S
# 累加最后一列所有状态的方案数 ans = sum(dp_prev) % MOD return ans