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) 所表示的边互不相同。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/count-subtrees-with-max-distance-between-cities

解答

/*
这个题，首先将由n个城市组成的城市集，用回溯法分成若干个子城市集。
其次，判断每个子城市集是否是联通的。
对于每一个子城市集，若它是联通的，则求出子城市集的最大距离（其实就是求“树的最大直径”），放入res数组中；
若子城市集不是联通的，则舍弃掉它。

对于"求出子城市集的最大距离"这个问题，有点类似于"求二叉树的最大直径"这个题；
不过这里的树不是二叉树，是一颗多叉树，不过思路类似。
*/
class Solution {
    boolean[][] con;
    Map<Integer, List<Integer>> map;
    Set<Integer> visited;
    Set<List<Integer>> subs;
    int maxDiameter;
    public void dfs(int start, int n, Stack<Integer> input){//用于求全部子城市集
        int size = input.size();
        if(size > 1 && check(input))subs.add(new ArrayList<Integer>(input));
        if(start == n)return;
        input.push(start);
        dfs(start+1, n, input);
        input.pop();

        dfs(start+1, n, input);
    }
    public boolean check(List<Integer> input){//用于判断子城市集是否联通
        int num=0, size=input.size();
        for(int i=0;i<size-1;i++)
            for(int j=i+1;j<size;j++)
                if(con[input.get(i)][input.get(j)])num++;//num表示边的数量
        return num == size-1;//n个节点的连通图，若有n-1条边，且无环，则该连通图为一棵树
    }
    public int dfsMax(int start, int n, List<Integer> input){//求每个子城市集的"最大距离"
        if(!visited.add(start))return 0;
        List<Integer> temp = map.get(start);
        PriorityQueue<Integer> thisLevel = new PriorityQueue<Integer>();//每一层维护一个最大堆
        thisLevel.offer(0);
        thisLevel.offer(0);
        //if(temp == null)return 0;
        for(int i:temp){
            if(!visited.contains(i) && input.contains(i)){
                thisLevel.offer(dfsMax(i, n, input));
                thisLevel.poll();
            }
        }
        int small = thisLevel.poll();
        int big = thisLevel.poll();
        maxDiameter = Math.max(big+small, maxDiameter);
        return Math.max(big, small)+1;
    }
    public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
        if(n==0)return new int[0];
        con = new boolean[n][n];
        //for(int i=0;i<n;i++)con[i][i]=true;

        map = new HashMap<>();
        for(int[] edge:edges){
            con[edge[0]-1][edge[1]-1]=true;
            con[edge[1]-1][edge[0]-1]=true;
            if(!map.containsKey(edge[0]-1))map.put(edge[0]-1, new ArrayList<Integer>());
            map.get(edge[0]-1).add(edge[1]-1);
            if(!map.containsKey(edge[1]-1))map.put(edge[1]-1, new ArrayList<Integer>());
            map.get(edge[1]-1).add(edge[0]-1);
        }
        //floyd
        // for(int k=0;k<n;k++)
        //     for(int i=0;i<n;i++)
        //         for(int j=0;j<n;j++)
        //             if(con[i][k] && con[k][k])con[i][j]=true;
        
        subs = new HashSet<List<Integer>>();
        dfs(0, n, new Stack<Integer>());
        visited = new HashSet<>();
        int[] result = new int[n-1];
        for(List<Integer> sub:subs){
            maxDiameter = 0;
            dfsMax(sub.get(0), n, sub);
            //System.out.println(sub);
            result[maxDiameter-1]++;
            visited.clear();
        }
        //System.out.println(subs);
        return result;
    }
}