LeetCode每日一题(2023/3/12)


1617. 统计子树中城市之间最大距离

给你 n 个城市,编号为从 1 到 n 。同时给你一个大小为 n-1 的数组 edges ,其中 edges[i] = [ui, vi] 表示城市 ui 和 vi 之间有一条双向边。题目保证任意城市之间只有唯一的一条路径。换句话说,所有城市形成了一棵 树 。

一棵 子树 是城市的一个子集,且子集中任意城市之间可以通过子集中的其他城市和边到达。两个子树被认为不一样的条件是至少有一个城市在其中一棵子树中存在,但在另一棵子树中不存在。

对于 d 从 1 到 n-1 ,请你找到城市间 最大距离 恰好为 d 的所有子树数目。

请你返回一个大小为 n-1 的数组,其中第 d 个元素(下标从 1 开始)是城市间 最大距离 恰好等于 d 的子树数目。

请注意,两个城市间距离定义为它们之间需要经过的边的数目。

示例 1:
输入:n = 4, edges = [[1,2],[2,3],[2,4]]
输出:[3,4,0]
解释:
子树 {1,2}, {2,3} 和 {2,4} 最大距离都是 1 。
子树 {1,2,3}, {1,2,4}, {2,3,4} 和 {1,2,3,4} 最大距离都为 2 。
不存在城市间最大距离为 3 的子树。
示例 2:

输入:n = 2, edges = [[1,2]]
输出:[1]
示例 3:

输入:n = 3, edges = [[1,2],[2,3]]
输出:[2,1]
 
提示:

2 <= n <= 15
edges.length == n-1
edges[i].length == 2
1 <= ui, vi <= n
题目保证 (ui, vi) 所表示的边互不相同。

解答:需要注意的是,枚举的子集不一定是一棵树,可能是森林(多棵树,多个连通块)。

我们可以在计算树形 DP 的同时去统计访问过的点,看看是否与子集相等,只有相等才是一棵树。

class Solution {
    // 邻接表表示图
    private List<Integer>[] g;
    // inSet[i] 表示节点 i 是否被选中,vis[i] 表示节点 i 是否被访问过
    private boolean[] inSet, vis;
    // ans[i] 表示直径为 i+1 的子图个数
    private int[] ans;
    // n 表示节点数,diameter 表示当前子图的直径
    private int n, diameter;

    public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
        this.n = n;
        g = new ArrayList[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int x = e[0] - 1, y = e[1] - 1; // 编号改为从 0 开始
            g[x].add(y);
            g[y].add(x); // 建树
        }

        ans = new int[n - 1];
        inSet = new boolean[n];
        f(0);
        return ans;
    }

    private void f(int i) {
        if (i == n) {
            // 枚举每个点作为子图的起点
            for (int v = 0; v < n; ++v)
                if (inSet[v]) {
                    vis = new boolean[n];
                    diameter = 0;
                    dfs(v);
                    break;
                }
            // 如果当前子图不连通,或者当前子图包含的节点不等于已选中的节点,则不计入答案
            if (diameter > 0 && Arrays.equals(vis, inSet))
                ++ans[diameter - 1];
            return;
        }

        // 不选节点 i
        f(i + 1);

        // 选节点 i
        inSet[i] = true;
        f(i + 1);
        inSet[i] = false; // 恢复现场
    }

    // 求以 x 为根节点的子树的直径
    private int dfs(int x) {
        vis[x] = true;
        int maxLen = 0;
        for (int y : g[x])
            if (!vis[y] && inSet[y]) {
                int ml = dfs(y) + 1;
                diameter = Math.max(diameter, maxLen + ml);
                maxLen = Math.max(maxLen, ml);
            }
        return maxLen;
    }
}

  目录