LeetCode 310: Minimum Height Trees

Question

Posting my code for a LeetCode problem, if you'd like to review, please do so. Thank you for your time!

Problem

For an undirected graph with tree characteristics, we can choose any node as the root. The result graph is then a rooted tree. Among all possible rooted trees, those with minimum height are called minimum height trees (MHTs). Given such a graph, write a function to find all the MHTs and return a list of their root labels.

Format

The graph contains n nodes which are labeled from 0 to n - 1. You will be given the number n and a list of undirected edges (each edge is a pair of labels).

You can assume that no duplicate edges will appear in edges. Since all edges are undirected, [0, 1] is the same as [1, 0] and thus will not appear together in edges.

Example 1 :

Input: n = 4, edges = [[1, 0], [1, 2], [1, 3]]

        0
        |
        1
       / \
      2   3 

Output: [1]

Example 2 :

Input: n = 6, edges = [[0, 3], [1, 3], [2, 3], [4, 3], [5, 4]]

     0  1  2
      \ | /
        3
        |
        4
        |
        5 

Output: [3, 4]

Note:

According to the definition of tree on Wikipedia: “a tree is an undirected graph in which any two vertices are connected by exactly one path. In other words, any connected graph without simple cycles is a tree.” The height of a rooted tree is the number of edges on the longest downward path between the root and a leaf.

Code

// The following block might slightly improve the execution time;
// Can be removed;
static const auto __optimize__ = []() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(NULL);
    std::cout.tie(NULL);
    return 0;
}();


// Most of headers are already included;
// Can be removed;
#include <cstdint>
#include <vector>
#include <unordered_set>
#include <algorithm>

static const struct Solution {
        using ValueType = std::uint_fast16_t;

        static const std::vector<int> findMinHeightTrees(
            const int n,
            const std::vector<std::vector<int>>& edges
        ) {
            std::vector<int> buff_a;
            std::vector<int> buff_b;
            std::vector<int>* ptr_a = &buff_a;
            std::vector<int>* ptr_b = &buff_b;

            if (n == 1) {
                buff_a.emplace_back(0);
                return buff_a;
            }

            if (n == 2) {
                buff_a.emplace_back(0);
                buff_a.emplace_back(1);
                return buff_a;
            }

            std::vector<Node> graph(n);

            for (const auto& edge : edges) {
                graph[edge[0]].neighbors.insert(edge[1]);
                graph[edge[1]].neighbors.insert(edge[0]);
            }

            for (ValueType node = 0; node < n; ++node) {
                if (graph[node].isLeaf()) {
                    ptr_a->emplace_back(node);
                }
            }

            while (true) {
                for (const auto& leaf : *ptr_a) {
                    for (const auto& node : graph[leaf].neighbors) {
                        graph[node].neighbors.erase(leaf);

                        if (graph[node].isLeaf()) {
                            ptr_b->emplace_back(node);
                        }
                    }
                }

                if (ptr_b->empty()) {
                    return *ptr_a;
                }

                ptr_a->clear();
                std::swap(ptr_a, ptr_b);
            }
        }

    private:
        static const struct Node {
            std::unordered_set<ValueType> neighbors;
            const bool isLeaf() {
                return std::size(neighbors) == 1;
            }
        };
};

References

• Problem

• Discuss

• Wiki

Answer 1

A triply nested loops always look scary. Especially with while (true).

First thing first, trust yourself. The leaf node (and the ptr_a contains only leafs) has only one neighbour. The

for (const auto& node : graph[leaf].neighbors)

loop is effectively

auto& node = graph[leaf].neighbors.begin();

Second, no naked loops please. And more functions please. The

for (const auto& leaf : *ptr_a)

prunes leaves from the tree. Factor it out into a prune_leaves function, which would return a set (technically a vector) of the new leaves:

leaves = prune_leaves(leaves, graph);

Third, the outer loop shall naturally terminate when less than 3 leaves remain.

Finally, separate IO from the business logic. That said, a code along the lines of

    graph = build_graph();
    leaves = collect_leaves(graph);
    while (leaves.size() < 3) {
        leaves = prune_leaves(leaves, graph);
    }
    return leaves;

would win my endorsement. Notice how ptr_a and ptr_b - which are not the most descriptive names - disappear.