Count Subtrees With Max Distance Between Cities

This page explains Java solution to problem Count Subtrees With Max Distance Between Cities using Floyd Warshall algorithm.

Problem Statement

There are n cities numbered from 1 to n. You are given an array edges of size n-1, where edges[i] = [ui, vi] represents a bidirectional edge between cities ui and vi. There exists a unique path between each pair of cities. In other words, the cities form a tree.

A subtree is a subset of cities where every city is reachable from every other city in the subset, where the path between each pair passes through only the cities from the subset. Two subtrees are different if there is a city in one subtree that is not present in the other.

For each d from 1 to n - 1, find the number of subtrees in which the maximum distance between any two cities in the subtree is equal to d.

Return an array of size n - 1 where the dth element (1-indexed) is the number of subtrees in which the maximum distance between any two cities is equal to d.

Notice that the distance between the two cities is the number of edges in the path between them.

Input: n = 4, edges = [[1,2],[2,3],[2,4]]
Output: [3,4,0]
Explanation: The subtrees with subsets {1,2}, {2,3} and {2,4} have a max distance of 1.
The subtrees with subsets {1,2,3}, {1,2,4}, {2,3,4} and {1,2,3,4} have a max distance of 2.
No subtree has two nodes where the max distance between them is 3.

Input: n = 2, edges = [[1,2]]
Output: [1]

Input: n = 3, edges = [[1,2],[2,3]]
Output: [2,1]

Solution

package com.vc.hard;

class CountSubtreesWithMaxDistanceBetweenCities {
    public int[] countSubgraphsForEachDiameter(int n, int[][] edges) {
        int[][] distance = new int[n][n];

        //Assign not reachable for all the combinations except where from == to
        for(int from = 0; from < n; from++) {
            for(int to = 0; to < n; to++) {
                if(from != to) distance[from][to] = Integer.MAX_VALUE;
            }
        }

        //If edge is present assign distance as 1
        for(int[] edge: edges) {
            int from = edge[0] - 1;
            int to = edge[1] - 1;
            distance[from][to] = 1;
            distance[to][from] = 1;
        }

        //Calculate distance using floyd Warshall
        for(int via = 0; via < n; via++) {
            for(int from = 0; from < n; from++) {
                for(int to = 0; to < n; to++) {
                    if(from != to &&
                            distance[from][via] != Integer.MAX_VALUE && distance[via][to] != Integer.MAX_VALUE) {
                        distance[from][to] = Math.min(distance[from][to],
                                distance[from][via] + distance[via][to]);
                    }
                }
            }
        }

        int[] res = new int[n - 1];
        //Loop through all the subsets and see if it is sub-tree
        //If it is subtree, return maxDistance between any two nodes in a sub tree
        for(int state = 0; state < (1 << n); state++) {
            int maxDistance = maxDistance(state, distance, n);
            if(maxDistance > 0) res[maxDistance - 1]++;
        }
        return res;
    }

    private int maxDistance(int state, int[][] distance, int n) {
        int countCity = 0, countEdge = 0, maxDistance = 0;
        for(int from = 0; from < n; from++) {
            if((state & (1 << from)) == 0) continue;
            countCity++;
            for(int to = from; to < n; to++) {
                if((state & (1 << to)) == 0) continue;
                countEdge += distance[from][to] == 1 ? 1 : 0;
                maxDistance = Math.max(maxDistance, distance[from][to]);
            }
        }
        if(countEdge != countCity - 1) return 0; //Subset is not valid subtree
        return maxDistance;
    }
}

Time Complexity

O(2^N) Where
N is total number of cities

Space Complexity

O(N²) Where
N is total number of cities