每日算法 — 使用java实现状态机DP：房屋染色最优策略与多状态转移设计

可莉爸爸 AI算法

引言

动态规划（DP）的核心在于状态的设计与转移。当问题中存在多种互斥的”模式”或”阶段”需要交替选择时，状态机动态规划（State Machine DP） 是一种强大的建模工具。本文以经典的房屋染色问题为载体，用 Java 完整实现状态机 DP 的设计思路、状态转移方程推导以及空间优化技巧，帮助读者掌握这一高阶 DP 技巧。

问题定义

房屋染色问题

有 $n$ 栋房屋排成一排，每栋房屋可以染成红色、蓝色或绿色中的一种。染色需要满足以下约束：
– 任意两栋相邻的房屋不能染成相同的颜色
– 染成红色花费 $r_i$，蓝色花费 $b_i$，绿色花费 $g_i$
– 求染完所有房屋的最小总花费

状态机视角

将每种颜色视为一个”状态”，整个染色过程就是在三种状态之间的转移过程：
– 当前房屋染红色 → 下一房屋只能染蓝色或绿色
– 当前房屋染蓝色 → 下一房屋只能染红色或绿色
– 当前房屋染绿色 → 下一房屋只能染红色或蓝色

这是一个典型的三状态有限状态机，状态之间的转移由相邻约束决定。

核心原理

状态定义

设 $dp[i][c]$ 表示将前 $i$ 栋房屋染完，且第 $i$ 栋染成颜色 $c$ 的最小花费。

其中 $c \in {0, 1, 2}$，分别代表红、蓝、绿三种颜色。

状态转移方程

对于第 $i$ 栋房屋，如果染成颜色 $c$，则第 $i-1$ 栋必须染成其他颜色：

$$dp[i][c] = cost[i][c] + \min_{c’ \neq c} {dp[i-1][c’]}$$

即：当前颜色花费 + 前一栋其他颜色中的最小花费。

初始化

$$dp[0][c] = cost[0][c], \quad c \in {0, 1, 2}$$

答案

$$\min_{c} {dp[n-1][c]}$$

基础实现

import java.util.*;

/**
 * 状态机DP：房屋染色问题基础实现
 */
public class HousePaintingDP {

    // 颜色常量
    public static final int RED = 0;
    public static final int BLUE = 1;
    public static final int GREEN = 2;
    public static final int COLOR_COUNT = 3;

    /**
     * 基础版本：O(n * k^2)，k为颜色数
     * @param costs costs[i][c] 表示第i栋染颜色c的花费
     * @return 最小总花费
     */
    public static int minCost(int[][] costs) {
        if (costs == null || costs.length == 0) return 0;

        int n = costs.length;
        int k = costs[0].length; // 颜色数量

        // dp[i][c] = 前i栋，第i栋染颜色c的最小花费
        int[][] dp = new int[n][k];

        // 初始化第一栋
        for (int c = 0; c < k; c++) {
            dp[0][c] = costs[0][c];
        }

        // 状态转移
        for (int i = 1; i < n; i++) {
            for (int c = 0; c < k; c++) {
                dp[i][c] = Integer.MAX_VALUE;
                // 前一栋必须染不同颜色
                for (int prevC = 0; prevC < k; prevC++) {
                    if (prevC != c) {
                        dp[i][c] = Math.min(dp[i][c], dp[i-1][prevC] + costs[i][c]);
                    }
                }
            }
        }

        // 答案取最后一栋的最小值
        int result = Integer.MAX_VALUE;
        for (int c = 0; c < k; c++) {
            result = Math.min(result, dp[n-1][c]);
        }
        return result;
    }

    /**
     * 空间优化版本：O(k) 空间
     * 由于 dp[i] 只依赖于 dp[i-1]，可以用滚动数组
     */
    public static int minCostSpaceOptimized(int[][] costs) {
        if (costs == null || costs.length == 0) return 0;

        int n = costs.length;
        int k = costs[0].length;

        int[] prev = new int[k];
        int[] curr = new int[k];

        // 初始化
        System.arraycopy(costs[0], 0, prev, 0, k);

        for (int i = 1; i < n; i++) {
            for (int c = 0; c < k; c++) {
                curr[c] = Integer.MAX_VALUE;
                for (int prevC = 0; prevC < k; prevC++) {
                    if (prevC != c) {
                        curr[c] = Math.min(curr[c], prev[prevC] + costs[i][c]);
                    }
                }
            }
            // 交换数组
            int[] temp = prev;
            prev = curr;
            curr = temp;
        }

        int result = Integer.MAX_VALUE;
        for (int c = 0; c < k; c++) {
            result = Math.min(result, prev[c]);
        }
        return result;
    }

    /**
     * 时间优化版本：O(n * k)
     * 利用前缀最小值和后缀最小值，避免内层循环
     */
    public static int minCostTimeOptimized(int[][] costs) {
        if (costs == null || costs.length == 0) return 0;

        int n = costs.length;
        int k = costs[0].length;

        int[] prev = new int[k];
        System.arraycopy(costs[0], 0, prev, 0, k);

        for (int i = 1; i < n; i++) {
            // 计算前缀最小值
            int[] leftMin = new int[k];
            leftMin[0] = Integer.MAX_VALUE; // 第一个位置没有左侧
            for (int j = 1; j < k; j++) {
                leftMin[j] = Math.min(leftMin[j-1], prev[j-1]);
            }

            // 计算后缀最小值
            int[] rightMin = new int[k];
            rightMin[k-1] = Integer.MAX_VALUE; // 最后一个位置没有右侧
            for (int j = k - 2; j >= 0; j--) {
                rightMin[j] = Math.min(rightMin[j+1], prev[j+1]);
            }

            // 当前状态 = costs[i][c] + min(左侧最小, 右侧最小)
            int[] curr = new int[k];
            for (int c = 0; c < k; c++) {
                curr[c] = costs[i][c] + Math.min(leftMin[c], rightMin[c]);
            }
            prev = curr;
        }

        int result = Integer.MAX_VALUE;
        for (int c = 0; c < k; c++) {
            result = Math.min(result, prev[c]);
        }
        return result;
    }
}

扩展：带额外约束的状态机 DP

问题升级：相邻两栋不能同色，且红色不能连续出现超过 2 次

这引入了”记忆”机制，需要记录连续红色的次数。状态扩展为 $dp[i][c][r]$，其中 $r$ 表示以第 $i$ 栋结尾连续红色的次数（0、1、2）。

/**
 * 增强版：红色不能连续超过 maxConsecutiveRed 次
 */
public class AdvancedHousePaintingDP {

    public static final int RED = 0;
    public static final int BLUE = 1;
    public static final int GREEN = 2;

    /**
     * 状态机DP：限制红色连续次数
     * @param costs 花费矩阵
     * @param maxConsecutiveRed 红色最大连续次数
     * @return 最小总花费
     */
    public static int minCostWithRedLimit(int[][] costs, int maxConsecutiveRed) {
        if (costs == null || costs.length == 0) return 0;

        int n = costs.length;
        int k = costs[0].length; // 假设 0=红, 1=蓝, 2=绿

        // dp[c][r] = 当前栋染颜色c，且连续红色次数为r的最小花费
        // r的取值：0（当前非红），1（当前红且前一栋非红），2（当前红且前一栋红）...
        int[][] dp = new int[k][maxConsecutiveRed + 1];
        for (int[] row : dp) Arrays.fill(row, Integer.MAX_VALUE / 2);

        // 初始化第一栋
        dp[RED][1] = costs[0][RED];     // 第一栋染红，连续1次
        dp[BLUE][0] = costs[0][BLUE];   // 第一栋染蓝，连续0次
        dp[GREEN][0] = costs[0][GREEN]; // 第一栋染绿，连续0次

        for (int i = 1; i < n; i++) {
            int[][] newDp = new int[k][maxConsecutiveRed + 1];
            for (int[] row : newDp) Arrays.fill(row, Integer.MAX_VALUE / 2);

            // 当前染红色
            for (int prevC = 0; prevC < k; prevC++) {
                if (prevC == RED) {
                    // 前一栋也是红色，连续次数+1
                    for (int r = 1; r < maxConsecutiveRed; r++) {
                        if (dp[prevC][r] < Integer.MAX_VALUE / 2) {
                            newDp[RED][r + 1] = Math.min(newDp[RED][r + 1], 
                                                         dp[prevC][r] + costs[i][RED]);
                        }
                    }
                } else {
                    // 前一栋非红，连续次数重置为1
                    for (int r = 0; r <= maxConsecutiveRed; r++) {
                        if (dp[prevC][r] < Integer.MAX_VALUE / 2) {
                            newDp[RED][1] = Math.min(newDp[RED][1], 
                                                     dp[prevC][r] + costs[i][RED]);
                        }
                    }
                }
            }

            // 当前染蓝色或绿色（非红色）
            for (int currC = 1; currC < k; currC++) {
                for (int prevC = 0; prevC < k; prevC++) {
                    if (prevC == currC) continue; // 相邻不能同色
                    for (int r = 0; r <= maxConsecutiveRed; r++) {
                        if (dp[prevC][r] < Integer.MAX_VALUE / 2) {
                            newDp[currC][0] = Math.min(newDp[currC][0], 
                                                       dp[prevC][r] + costs[i][currC]);
                        }
                    }
                }
            }

            dp = newDp;
        }

        int result = Integer.MAX_VALUE;
        for (int c = 0; c < k; c++) {
            for (int r = 0; r <= maxConsecutiveRed; r++) {
                result = Math.min(result, dp[c][r]);
            }
        }
        return result;
    }
}

通用状态机 DP 框架

/**
 * 通用状态机DP框架
 * @param <State> 状态类型
 */
public class StateMachineDPFramework<State> {

    /**
     * 状态转移接口
     */
    @FunctionalInterface
    public interface Transition<State> {
        List<State> getNextStates(State current);
    }

    /**
     * 通用状态机DP求解
     * @param n 步骤数
     * @param states 所有可能的状态集合
     * @param costFunc cost(i, state) = 第i步选择state的代价
     * @param transition 状态转移规则
     * @return 最小总代价
     */
    public static <State> long solve(
            int n,
            List<State> states,
            CostFunction<State> costFunc,
            Transition<State> transition) {

        Map<State, Long> dp = new HashMap<>();

        // 初始化
        for (State s : states) {
            dp.put(s, costFunc.cost(0, s));
        }

        // 递推
        for (int i = 1; i < n; i++) {
            Map<State, Long> newDp = new HashMap<>();
            for (State curr : states) {
                long minCost = Long.MAX_VALUE;
                for (State prev : states) {
                    if (transition.getNextStates(prev).contains(curr)) {
                        minCost = Math.min(minCost, dp.get(prev) + costFunc.cost(i, curr));
                    }
                }
                newDp.put(curr, minCost);
            }
            dp = newDp;
        }

        return dp.values().stream().mapToLong(Long::longValue).min().orElse(0L);
    }

    @FunctionalInterface
    public interface CostFunction<State> {
        long cost(int step, State state);
    }
}

测试用例

public class HousePaintingTest {

    public static void main(String[] args) {
        // 测试用例 1：3栋房屋，3种颜色
        int[][] costs1 = {
            {17, 2, 17},   // 第0栋：红17，蓝2，绿17
            {16, 16, 5},   // 第1栋：红16，蓝16，绿5
            {14, 3, 19}    // 第2栋：红14，蓝3，绿19
        };

        System.out.println("测试1 - 基础版本: " + HousePaintingDP.minCost(costs1));
        System.out.println("测试1 - 空间优化: " + HousePaintingDP.minCostSpaceOptimized(costs1));
        System.out.println("测试1 - 时间优化: " + HousePaintingDP.minCostTimeOptimized(costs1));
        // 最优解：蓝(2) + 绿(5) + 蓝(3) = 10

        // 测试用例 2：更多颜色
        int[][] costs2 = {
            {1, 5, 3, 2},
            {4, 2, 6, 1},
            {3, 1, 4, 5},
            {2, 4, 1, 3}
        };
        System.out.println("\n测试2 - 4种颜色: " + HousePaintingDP.minCost(costs2));

        // 测试用例 3：大规模数据
        int[][] costs3 = generateRandomCosts(100000, 3);
        long start = System.currentTimeMillis();
        int result3 = HousePaintingDP.minCostTimeOptimized(costs3);
        System.out.println("\n测试3 - 大规模(n=100000, k=3): " + result3);
        System.out.println("  耗时: " + (System.currentTimeMillis() - start) + "ms");

        // 测试用例 4：限制红色连续次数
        int[][] costs4 = {
            {5, 2, 3},
            {4, 5, 1},
            {3, 2, 4},
            {5, 1, 3},
            {2, 4, 5}
        };
        System.out.println("\n测试4 - 红色最多连续2次: " + 
            AdvancedHousePaintingDP.minCostWithRedLimit(costs4, 2));
    }

    static int[][] generateRandomCosts(int n, int k) {
        Random rand = new Random(42);
        int[][] costs = new int[n][k];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < k; j++) {
                costs[i][j] = rand.nextInt(100) + 1;
            }
        }
        return costs;
    }
}

复杂度分析

版本 时间复杂度 空间复杂度 适用场景
基础版 $O(n \cdot k^2)$ $O(n \cdot k)$ 理解状态转移逻辑
空间优化 $O(n \cdot k^2)$ $O(k)$ 内存受限
时间优化 $O(n \cdot k)$ $O(k)$ 颜色数较多（$k > 10$）
限制连续版 $O(n \cdot k^2 \cdot m)$ $O(k \cdot m)$ 含连续次数约束

其中 $n$ 为房屋数，$k$ 为颜色数，$m$ 为最大连续次数。

状态机 DP 的应用场景

状态机 DP 的适用范围远不止房屋染色：

  • 股票买卖问题：持有/不持有两种状态，含冷却期或交易次数限制
  • 字符串模式匹配：正则表达式引擎的状态机实现
  • 游戏 AI：有限状态机驱动的 NPC 行为决策
  • 编译器词法分析：DFA/NFA 的状态转移表构建
  • 工作流引擎：流程节点的状态流转与路径优化

总结

状态机 DP 的本质是将问题的约束条件建模为状态之间的转移规则。通过本文的房屋染色问题，我们学习了：

  1. 基础状态机：三状态互斥转移，$O(n \cdot k^2)$ 解法
  2. 空间优化：滚动数组将空间降至 $O(k)$
  3. 时间优化：前缀/后缀最小值将时间降至 $O(n \cdot k)$
  4. 状态扩展：引入”连续计数”维度处理更复杂的约束

掌握状态机 DP 的设计范式，能够帮助你优雅地解决一类”多模式交替选择”的优化问题，是算法竞赛和工程实践中的利器。

Tagged with: