Problem Description

With time, it gets a bit tricky to find truly interesting challenges on LeetCode, but a couple of days ago I came across one. I found that this hard LeetCode challenge “1617. Count Subtrees With Max Distance Between Cities” isn’t actually that hard and at the same time is very informative from learning perspective: it brings together a few mid-level ideas, and that’s exactly why I like it.

In short, the challenge is to count, for every possible diameter d, how many subtrees of a given tree have d as a diameter.

Basic idea

Given the problem constraint that the number of vertices n≤15, we can safely assume that enumerating all subsets of vertices is feasible, as there are only 2^15=32768 possible subsets. Of course to enumerate over subsets one can apply the bit mask approach.

Now, for each subset of vertices, we can iterate over all the edges in the original graph and check if we can form a valid subtree using these vertices and edges. Since we cannot introduce new cycles(otherwise we get a cycle in the original graph), the condition for a valid tree is simple: the number of vertices should be equal to the number of edges plus one.

Finally for each sub-tree we need to compute its diameter and add this to the resulting counting array.

Diameter of a tree using Two-Pass BFS

Given a tree, we need to find its diameter, i.e., the maximum distance between any pair of vertices. This can be done in several ways, such as running the classical Dijkstra algorithm from every node, or using Floyd-Warshall, the later one for instance requires roughly |V|^3 operations where |V| is the amount of vertices we have. However, the reason I’m writing this post is that a couple of years ago, I learned about a beautiful alternative: the Two-Pass BFS algorithm. This method works with just two traversals of the graph! I first came across it in this video link, and since then, I’ve been eagerly waiting for a good challenge to apply it. By the way, the video is based on explanations from the beautiful book Explaining Algorithms Using Metaphors. Someday, I might write about other algorithms from there, so I highly recommend this book to fellow engineers.

As the name suggests, the Two-Pass BFS algorithm works with two BFS passes over the graph:

1. First, we perform a BFS starting from any (randomly selected)vertex to find a vertex A that is at the maximum distance from the starting point.
2. Next, we run the BFS again, but this time starting from vertex A, and compute the maximum distance we can travel from A.

The claim is that the maximum distance found in the second pass is actually the diameter of the tree. If you’re hearing this for the first time, it might sound like a miracle, but once you visualize it — like in the video linked above — you can truly appreciate how elegant and beautiful this fact is!

Quick analysis

Let’s analyze the total time complexity: We start with n vertices and n−1 edges. We need to loop over 2^n subsets. For each subset, we perform n operations to add vertices from a bit mask, n−1 operations to check edges, and then we need to run BFS twice, which is also bounded by 2n. So, in total, the time complexity is O(n⋅2^n), which is quite reasonable and might work.

The Code

A quick implementation I came up with:

public class Solution {
    public int[] CountSubgraphsForEachDiameter(int n, int[][] edges) 
    {
        var result = new int[n - 1];
        for(int mask = 1; mask < (1 << n); ++mask)
        {
            var graph = new Dictionary<int, HashSet<int>>();
            var edgesCount = 0;

            for(int i = 0; i < n; ++i)
            {
                if((mask & (1 << i)) > 0)
                {
                    graph.TryAdd(i, new HashSet<int>());
                }
            }

            foreach(var edge in edges)
            {
                var start = edge[0] - 1;
                var end = edge[1] - 1;
                if( graph.ContainsKey(start) && 
                    graph.ContainsKey(end) )
                {
                    graph[start].Add(end);
                    graph[end].Add(start);
                    ++edgesCount;

                }
            }

            if(graph.Keys.Count - 1 == edgesCount && edgesCount > 0)
            {
                var (first, maxDistance) = BFS(graph, BFS(graph, graph.Keys.First()).Item1);

                if(maxDistance > 0 && maxDistance <= n - 1)
                {
                    result[maxDistance - 1] += 1;        
                }
            }
        }

        return result;
    }

    public (int, int) BFS(Dictionary<int, HashSet<int>> graph, int start)
    {
        var toVisit = new HashSet<int>();
        toVisit.Add(start);
        var visited = new HashSet<int>();
        
        var count = 0;
        var lastIndex = start;

        while(toVisit.Count > 0)
        {
            ++count;
            var newToVisit = new HashSet<int>();
            foreach(var candidate in toVisit)
            {
                visited.Add(candidate);
                if(graph.ContainsKey(candidate))
                {
                    foreach(var newCandidate in graph[candidate])
                    {
                        if(!visited.Contains(newCandidate))
                        {
                            newToVisit.Add(newCandidate);
                            lastIndex = newCandidate;
                        }
                    }
                }
            }
            toVisit = newToVisit;
        }

        return (lastIndex, count - 1);
    } 
}

And surprisingly, this passes all the tests: