方法探索｜一场树形 DP 深度游

公众号新闻

2023-02-23 18:02

很多人会问刷多少道 DP 才能培养出 DP 感？

其实答案就是多做，日积月累，多做的深层含义就是理解，可以先刷二叉树类的题目，又或者看看下面这位站内同学的分享，看看有没有启发。

解题思路

👉 834.树中距离之和

其实这位同学在之前的第 328 场周赛中，也遇到过此类困难的题型。大家深究其实可以发现，可以用动态规划来解题，但如果想不到怎么做状态转移的话，可以去看下 👉 官方题解。通过官方解题思路，我们可以了解到下图中的状态转移方程：

而不是

树状 DP 的特点，就是动态规划的每个子问题，一一对应于一颗子树。可以从由单个叶子节点构成的子树开始，自底向上进行规划。也可以从根节点开始，自上而下地、使用带缓存的方式进行规划。后者本质就是一个对整棵树 DFS 的过程。

树状 DP，时间复杂度为 O(n)。

那么，有一类与无根树相关的题目，根节点可以任意选择，需要对每一种可能的根节点，都运行一次树状 DP，所以整体时间复杂度为 O(n*2)。可以通过反复换根的方式，将这种题目的时间复杂度度优化到 O(n)。这个换根的过程，也采取 DFS 的方式。

上面这道题其实就属于这种题型，需要先随意选择一个根，运行一次树状 DP，然后额外运行一次 DFS 进行换根。也就是总共进行两次 DFS。之前也会看到许多扣友将这种两次 DFS 的思路，叫做换根 DP。

编码细节

理解了思路，该同学开始尝试写代码，又发现了更多细节。

在阐述中发现树状 DP，如果采用递归的方式会比较好写，如果用自底向上的方式，需要先进行层序遍历（BFS），对节点按层分组，然后从最底层开始进行动态规划。

为了循序渐进，可以先写了一个进行 N 次树状 DP 的暴力版本，这样做显然会超时，但该同学是为了分步进行练习，先找到写树状 DP 的感觉。毕竟一口真的很难吃成个胖子。

再确保暴力版本正确后，开始尝试换根，这个过程调试起来很痛苦，时长又久。所以该同学得出：“在思考解题的过程中一定要时刻记住换根的本质”，防止思维混乱：子节点的子节点多了一个，父节点的子节点少了一个。换根的过程，虽然也是 DFS，但和动态规划毫无关系，大家不要写混了。

代码实践

暴力版本：N 次树状 DP

*暴力版会超时，仅做为练习。

代码

class Solution {    public int[] sumOfDistancesInTree(int n, int[][] edges) {        List<Integer>[] adt = new List[n];        for(int i = 0; i < n; i++)adt[i] = new ArrayList<Integer>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }        int[] ret = new int[n];        for(int i = 0; i < n; i++){            //System.out.printf("%s--------------------\n", i);            // 以 i 为树的根节点，进行自上而下带缓存的的动态规划            int[] dp = new int[n]; // dp[i] 代表：对于以节点 i 为根的子树，i 到所有后代的距离的和            int[] size = new int[n]; // size[i] 代表以节点i 为根的子树中，所有节点的数量            helper(n, adt, i, i, dp, size);            ret[i] = dp[i];        }        return ret;    }
    /**    用动态规划，递归地求出以 u 为根节的子树中，u 到所有后代节点的距离之和。pa 代表 u 的直接父节点    调用一次 helper，会求出 dp[u] 和 size[u]     */    private void helper(int n, List<Integer>[] adt, int pa, int u, int[] dp, int[] size){        if(size[u] > 0)return;
        size[u]++;        for(int v: adt[u]){            if(v == pa)continue;            helper(n, adt, u, v, dp, size);            dp[u] += dp[v] + size[v];             size[u] += size[v];        }        //System.out.printf("%s: %s, %s\n", u, size[u], dp[u]);    }}

树状 DP + DFS 换根

换根的逻辑，被单独提取成了 SWAP 方法，可以让代码简洁一点：

代码

class Solution {    public int[] sumOfDistancesInTree(int n, int[][] edges) {        List<Integer>[] adt = new List[n];        for(int i = 0; i < n; i++)adt[i] = new ArrayList<Integer>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }        int[] ret = new int[n];        int[] dp = new int[n]; // dp[i] 代表以 i 为根的子树中， i 到所有后代节点的距离之和。i 的父节点可以有多种情况        int[] size = new int[n]; // size[i] 代表以节点 i 为根的【子树】中，所有节点的数量                helper(n, adt, -1, 0, dp, size);        ret[0] = dp[0];        dfs(adt, -1, 0, ret, dp, size);        return ret;    }
    /**    用动态规划，递归地求出以 u 为根节的子树中，u 到所有后代节点的距离之和。pa 代表 u 的直接父节点    调用一次 helper，会求出 dp[u] 和 size[u]     */    private void helper(int n, List<Integer>[] adt, int pa, int u, int[] dp, int[] size){        if(size[u] > 0)return;
        size[u]++;        for(int v: adt[u]){            if(v == pa)continue;            helper(n, adt, u, v, dp, size);            dp[u] += dp[v] + size[v];             size[u] += size[v];        }        //System.out.printf("%s: %s, %s\n", u, size[u], dp[u]);    }
    // 用回溯法不断换根    private void dfs(List<Integer>[] adt,int pa, int u, int[] ret, int[] dp, int[] size){        for(int v: adt[u]){            if(v == pa)continue;            swap(u, v, dp, size);            ret[v] = dp[v];            //System.out.printf("%s %s, %s\n", u, v, ret[v]);            dfs(adt, u, v, ret, dp, size);            swap(v, u, dp, size);        }    }
    private void swap(int u, int v, int[] dp, int[] size){        dp[u] -= dp[v] + size[v];        size[u] -= size[v];        dp[v] += dp[u] + size[u];        size[v] += size[u];    }}

举一反三：更多树状 DP 题目（含代码）

👉 2538. 最大价值和与最小价值和的差值

如果这道题比较痛苦，可以跳过，先做后面的中等题 1245。

这道题有两点需要注意：

确保每次换根在 O(1) 内完成，需要点技巧：对要被替换的某子树根节点 u ，先遍历其邻接点一遍，求出 DP 值的最大之和次大值。如果不注意，会喜提超时，详见下面的错误 版本 1 和正确 版本 2。

也可以换一种思路，不用换根，直接进行一次树状 DP 即可（借鉴灵茶山艾府解法）。

下面的 4 个版本中，后 3 个都能良好的工作，版本 3 和 版本 4 思路是一样的。可以看到，虽然都用了树形 DP，但是定义的子问题有差别，所以状态转移的过程，也有差异。

（1）换根但 TLE 版代码

/*单次换根没有在 O(1)内完成，导致超时。834 题可以将换根提取到一个简单的 swap 方法中，比较简洁，但是这道题不可以。通过额外遍历一次邻接点，求出最大值和次大值，确保了O(1)的均摊单次换根代价。*//*单次换根没有在 O(1)内完成，导致超时。834 题可以将换根提取到一个简单的 swap 方法中，比较简洁，但是这道题不可以。通过额外遍历一次邻接点，求出最大值和次大值，确保了O(1)的均摊单次换根代价。*/class Solution {    private List<Integer>[] adt; // 邻接表    private long max = 0;    private long[] dp; // dp[i] 代表以 i 为根节点的子树中，i 最大的【开销】    private int[] price;
    public long maxOutput(int n, int[][] edges, int[] price) {        this.adt = new List[n];        this.max = 0;        this.dp = new long[n];        this.price = price;
        // 构造邻接表        for(int i = 0; i < n; i++)adt[i] = new ArrayList<>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        // 以节点 0 为根，运行一次自上而下的带缓存的树形动态规划        helper(-1, 0);        max = dp[0];        // 用回溯法不断换根，求出 max        dfs(-1, 0);        return max;    }
    // 自上而下的带缓存的树形动态规划    // pa代表 u 的父节点    private void helper(int pa, int u){        for(int v: adt[u]){            if(v == pa)continue;            helper(u, v);            dp[u] = Math.max(dp[u], dp[v] + price[v]);        }        System.out.printf("%s,%s\n", u, dp[u]);    }
    private void dfs(int pa, int u){        for(int v: adt[u]){            if(v == pa)continue;            swap(u, v);            max = Math.max(dp[v], max);            System.out.printf("换根：%s,%s, %s, %s, %s\n", u, v, dp[u], dp[v], max);            dfs(u, v);            swap(v, u);        }    }
    private void swap(int u, int v){        dp[u] = 0;        for(int x: adt[u]){            if(x == v)continue;            dp[u] = Math.max(dp[u], dp[x] + price[x]);        }        dp[v] = Math.max(dp[v], dp[u] + price[u]);    }}

（2）换根线性复杂度版代码

class Solution {    private List<Integer>[] adt; // 邻接表    private long max = 0;    private long[] dp; // dp[i] 代表以 i 为根节点的子树中，i 最大的【开销】    private long ret;    private int[] price;
    public long maxOutput(int n, int[][] edges, int[] price) {        this.adt = new List[n];        this.ret = 0;        this.dp = new long[n];        this.price = price;
        // 构造邻接表        for(int i = 0; i < n; i++)adt[i] = new ArrayList<>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        // 以节点 0 为根，运行一次自上而下的带缓存的树形动态规划        helper(-1, 0);        ret = dp[0];        // 用回溯法不断换根，求出 max        dfs(-1, 0);        return ret;    }
    // 自上而下的带缓存的树形动态规划    // pa代表 u 的父节点    private void helper(int pa, int u){        for(int v: adt[u]){            if(v == pa)continue;            helper(u, v);            dp[u] = Math.max(dp[u], dp[v] + price[v]);        }        // System.out.printf("%s,%s\n", u, dp[u]);    }
    private void dfs(int pa, int u){        long max = 0, max2 = 0; // 最大值和次大值        for(int v: adt[u]){
            long val = dp[v] + price[v];            if(val >= max){                max2 = max;                max = val;            }else{                max2 = Math.max(max2, val);            }        }        for(int v: adt[u]){            if(v == pa)continue;            long dpu = dp[u], dpv = dp[v];            if(dp[v] + price[v] == max)dp[u] = max2;            dp[v] = Math.max(dp[v], dp[u] + price[u]);            ret = Math.max(dp[v], ret);            // System.out.printf("换根：%s,%s, %s, %s, %s\n", u, v, dp[u], dp[v], ret);            dfs(u, v);            dp[u] = dpu;            dp[v] = dpv;        }    }
}

（3）不换根仿灵佬版代码

/*本代码参考了灵佬的版本，但是保留了 dp 数组，更方便我自己理解。灵佬原版中，将 dp 数组彻底优化掉了，更节省空间，但是更费脑细胞*/class Solution {    private List<Integer>[] adt; // 邻接表    private long max = 0;    private long[] dp; // dp[i] 代表以 i 为根的子树中，i 能达到的最大价值    private long[] dp2; // dp 2[i] 代表以 i 为根的子树中，在不累加叶子节点的价值是时，i 能达到的最大价值    private long ret;    private int[] price;
    public long maxOutput(int n, int[][] edges, int[] price) {        this.adt = new List[n];        this.ret = 0;        this.dp = new long[n];        this.dp2 = new long[n];        this.price = price;
        // 构造邻接表        for(int i = 0; i < n; i++)adt[i] = new ArrayList<>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        // 以节点 0 为根，运行一次自上而下的带缓存的树形动态规划        helper(-1, 0);        return max;    }
    // 自上而下的带缓存的树形动态规划    // pa代表 u 的父节点    private void helper(int pa, int u){        dp[u] = price[u];        for(int v: adt[u]){            if(v == pa)continue;            helper(u, v);            max = Math.max(Math.max(dp[u] + dp2[v], dp2[u] + dp[v]), max);            dp[u] = Math.max(dp[u], dp[v] + price[u]);            dp2[u] = Math.max(dp2[u], dp2[v] + price[u]);        }    }
}

（4）灵佬原版代码

/**本代码来着灵佬的题解：https://leetcode.cn/problems/difference-between-maximum-and-minimum-price-sum/solution/by-endlesscheng-5l70/灵佬的版本空间复杂度更低*/class Solution {    private List<Integer>[] g;    private int[] price;    private long ans;
    public long maxOutput(int n, int[][] edges, int[] price) {        this.price = price;        g = new ArrayList[n];        Arrays.setAll(g, e -> new ArrayList<>());        for (var e : edges) {            int x = e[0], y = e[1];            g[x].add(y);            g[y].add(x); // 建树        }        dfs(0, -1);        return ans;    }
    // 返回带叶子的最大路径和，不带叶子的最大路径和    private long[] dfs(int x, int fa) {        long p = price[x], maxS1 = p, maxS2 = 0;        for (var y : g[x])            if (y != fa) {                var res = dfs(y, x);                long s1 = res[0], s2 = res[1];                // 前面最大带叶子的路径和 + 当前不带叶子的路径和                // 前面最大不带叶子的路径和 + 当前带叶子的路径和                ans = Math.max(ans, Math.max(maxS1 + s2, maxS2 + s1));                maxS1 = Math.max(maxS1, s1 + p);                maxS2 = Math.max(maxS2, s2 + p); // 这里加上 p 是因为 x 必然不是叶子            }        return new long[]{maxS1, maxS2};    }}

中等题：

👉 1245.树的直径

这道题可以采用三种方式求解思路：

1、树形 DP + 换根，换根时，对任一颗子树的根节点 u，需要遍历其邻接点，并从中求出 DP 值最大和次大的子节点。这个思路类似于前面 2538. 最大价值和与最小价值和的差值 中的代码 版本 2。如果不这样做，在换根时遍历一遍邻接点，换根会超时。

2、一次树形 DP 即可。代价是做动态规划时，要同时维护每颗子树的根节点

能达到的最大路径和次大路径。所有子树中，最大路径和次大路径的和的最大值就是答案。

3、改进一下 思路 2，详见下面的代码 版本 3。

下面的三个版本，都能正常工作。但是状态转移的方式，各不相同。

（1）换根版代码

class Solution {    public int treeDiameter(int[][] edges) {        int n = edges.length + 1;        int[] dp = new int[n]; //  dp[i] 代表以 i 为根节点的子树中，i 到其后代节点的最长路径
        // 生成邻接表        List<Integer>[] adt = new List[n];        for(int i = 0; i < n; i++) adt[i] = new ArrayList<Integer>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        dfs(adt, -1, 0, dp);        int[] max = new int[]{dp[0]};        dfs2(adt, -1, 0, dp, max);        return max[0];    }
    private void dfs(List<Integer>[] adt, int pa, int u, int[] dp){        for(int v: adt[u]){            if(pa == v)continue;            dfs(adt, u, v, dp);            dp[u] = Math.max(dp[u], dp[v] + 1);        }        // System.out.printf("%s, %s\n", u, dp[u]);    }
    private void dfs2(List<Integer>[] adt, int pa, int u, int[] dp, int[] max){        int m1 = 0, m2 = 0; //最大值和次大值        for(int v: adt[u]){            int val = dp[v] + 1;            if(m1 <= val){                m2 = m1;                m1 = val;            }else{                m2 = Math.max(m2, val);            }        }        for(int v: adt[u]){            if(pa == v)continue;            int baku = dp[u], bakv = dp[v];            if(dp[v] + 1 == dp[u])dp[u] = m2;            dp[v] = Math.max(dp[v], dp[u] + 1);            max[0] = Math.max(max[0], dp[v]);            dfs2(adt, u, v, dp, max);            dp[u] = baku;            dp[v] = bakv;        }    }}

（2）不换根，一趟树形 DP版代码

class Solution {    public int treeDiameter(int[][] edges) {        int n = edges.length + 1;        int[] dp = new int[n]; //  dp[i] 代表以 i 为根节点的子树中，i 到其他后代节点的最长路径
        // 生成邻接表        List<Integer>[] adt = new List[n];        for(int i = 0; i < n; i++) adt[i] = new ArrayList<Integer>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        dfs(adt, -1, 0, dp);        int[] max = new int[]{dp[0]};        dfs2(adt, -1, 0, dp, max);        return max[0];    }
    private void dfs(List<Integer>[] adt, int pa, int u, int[] dp){        for(int v: adt[u]){            if(pa == v)continue;            dfs(adt, u, v, dp);            dp[u] = Math.max(dp[u], dp[v] + 1);        }        // System.out.printf("%s, %s\n", u, dp[u]);    }
    private void dfs2(List<Integer>[] adt, int pa, int u, int[] dp, int[] max){        int m1 = 0, m2 = 0; //最大值和次大值        for(int v: adt[u]){            int val = dp[v] + 1;            if(m1 <= val){                m2 = m1;                m1 = val;            }else{                m2 = Math.max(m2, val);            }        }        for(int v: adt[u]){            if(pa == v)continue;            int baku = dp[u], bakv = dp[v];            if(dp[v] + 1 == dp[u])dp[u] = m2;            dp[v] = Math.max(dp[v], dp[u] + 1);            max[0] = Math.max(max[0], dp[v]);            dfs2(adt, u, v, dp, max);            dp[u] = baku;            dp[v] = bakv;        }    }}

（3）不换根，一趟树形 DP版代码

class Solution {    public int treeDiameter(int[][] edges) {        int n = edges.length + 1;        int[] dp = new int[n]; //  dp[i] 代表以 i 为根节点的子树中，i 到其他后代节点的最长路径        // 生成邻接表        List<Integer>[] adt = new List[n];        for(int i = 0; i < n; i++) adt[i] = new ArrayList<Integer>();        for(int[] e: edges){            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        int[] max = new int[1];        dfs(adt, -1, 0, dp, max);        return max[0];    }
    /**    pa 是 u 的父节点     */    private void dfs(List<Integer>[] adt, int pa, int u, int[] dp, int[] max){        for(int v: adt[u]){            if(pa == v)continue;            dfs(adt, u, v, dp, max);            // 父节点当能达到的最长路径，加上当前子节点到达的最长路径，加上 1，然后更新 max            max[0] = Math.max(max[0], dp[u] + dp[v] + 1);            dp[u] = Math.max(dp[v] + 1, dp[u]);        }        // System.out.printf("%s, %s, %s\n", u, dp[u], dp2[u]);    }}

中等题：

👉 310.最小高度树

当你悟出了前面三道题之后，那这道题用树状 DP + 换根，思路也会变得自然，解题会很轻松。在换根时，依然采用了求最大和次大值的技巧。

这道题还有很多别的解法，可参考官方题解的 3 种解法，前两种比较数学化，而第三种会比较有趣。

最后奉上站内其他同学的解法，感兴趣的可以点击查看 👉「宫水三叶」树形 DP 运用题...

这位小伙伴是下面的思路：

树状 DP + 换根代码

class Solution {    private List<Integer>[] adt; // 邻接表    private int[] dp; // dp[i] 代表以 i 为根的子树的高度。注意，整棵树的根，是动态可变的。    private int[] dp2; // dp2 [i] 代表整棵树以 i 根时，树的高度    private int min;
    public List<Integer> findMinHeightTrees(int n, int[][] edges) {        dp = new int[n];        dp2 = new int[n];        min = Integer.MAX_VALUE;        adt = new List[n];        Arrays.setAll(adt, value -> new ArrayList<>());        for (int[] e : edges) {            adt[e[0]].add(e[1]);            adt[e[1]].add(e[0]);        }
        // 进行一次树状 dp，以 0 为根节点        helper(-1, 0);        dp2[0] = dp[0];        min = Math.min(dp[0], min);
        // 换根        dfs(-1, 0);
        List<Integer> ret = new ArrayList<>();        for(int i = 0; i < n; i++){            if(dp2[i] == min)ret.add(i);        }        return ret;    }
    /**     * 换根     * @param pa u 的父节点     * @param u     */    private void dfs(int pa, int u) {        int m1 = 0, m2 = 0; // 最大和次大        for (int v : adt[u]) {            int val = dp[v] + 1;            if(m1 <= val){                m2 = m1;                m1 = val;            }else{                m2 = Math.max(m2, val);            }        }        for (int v : adt[u]) {            if(v == pa)continue;            int baku = dp[u], bakv = dp[v];            if(dp[v] + 1 == dp[u])dp[u] = m2;            dp[v] = Math.max(dp[u] + 1, dp[v]); // 求出以 v 为整颗树的根时，树的高度
            min = Math.min(min, dp[v]);            dp2[v] = dp[v];
            // System.out.printf("换根%s, %s,%s, %s, %s\n",u, v, dp[v], min, ret);            dfs(u, v);            dp[u] = baku;            dp[v] = bakv;
        }    }
    /*    * 自上而下带缓存的动态规划，pa 是 u 的父节点    * */    private void helper(int pa, int u) {        for (int v : adt[u]) {            if(pa == v)continue;            helper(u,v);            dp[u] = Math.max(dp[u], dp[v] + 1);        }        //System.out.printf("%s, %s\n", u, dp[u]);    }}

写在最后：

解完这 4 道题，对树状 DP 感觉会没那么神秘和恐怖了。其实不管有没有“树状”这个修饰词，DP 最让人害怕的还是怎么想出状态转移方程。不同的状态转移思路，对应不同的解法。这四道题有个共同点，那就是都可以转换为求“最长路径”。也可能并不是所有的树状 DP 问题都会有这个共性，如果你还有更多的思路，欢迎投稿分享。

换根这个技巧，掌握了以后比较容易上手；追求不换根，反而会比较难。

BY /

本文作者：老宋

编辑&版式：Janson

点个在看，少个 bug

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。

来源: qq

点击查看作者最近其他文章