Catch the bug: diameter of a tree.

The following function to find the diameter (width) of a tree presents some nice propreties:

• (Almost) Representation agnostic.
  • The type of nodes (simple int used an index, string identifiant ...) doesn't matter.
  • The dict children might map the leaves to an empty sequence, or not contains them altogether.
• The code is short, efficient and could be used for competitive programming, if it wasn't bugged.

As an exercice, try to find the design flaw and the real bug hidden in this cute snippet.

import heapq


def tree_height_diameter(children, root):
    """ Return height and diameter of (sub)tree. """
    cs = children.get(root, [])
    if cs:
        # Collect and unzip.
        hs, ds = zip(*map(lambda x: tree_height_diameter(children, x), cs))
        hs = heapq.nlargest(2, hs)
        return hs[0] + 1, max(max(ds), sum(hs) + len(hs))
    else:
        return 0, 0  # Leaf

TODO: Explain the code, if popular demand.

view raw tree_diameter_bugged.md hosted with ❤ by GitHub