Implementation of breadth-first traversal for m-ary trees in Python and OCaml.

m_ary_tree.ml

type 'a t = Node of 'a * 'a t list

let get_value (Node (value, _)) = value
let get_children (Node (_, children)) = children

let traverse_breadth f tree =
  let rec traverse_breadth_acc queue acc =
    if List.is_empty queue
    then List.rev acc
    else
      traverse_breadth_acc
        (List.concat_map get_children queue)
        (f (List.rev_map get_value queue) acc)
  in
  traverse_breadth_acc [ tree ] []

let get_layers tree = traverse_breadth List.cons tree
let to_list_breadth tree = traverse_breadth ( @ ) tree

m_ary_tree.py

#!/usr/bin/env -S python -B

"""
Implementation of breadth-first traversal for m-ary trees.
"""

import itertools
import typing
import unittest
from collections.abc import Collection
from collections.abc import Generator
from itertools import chain

class Tree[A](typing.NamedTuple):
    """
    A m-ary tree.
    """

    value: A
    subtrees: list[Tree[A]]

def leaf[A](value: A) -> Tree[A]:
    """
    Return a single leaf, which is a tree without subtrees.
    """

    return Tree(value, [])

def get_layers[A](tree: Tree[A]) -> list[list[A]]:
    """
    Perform a breadth-first traversal, returning each layer as a sublist.
    """

    layers = []
    queue = [tree]

    while queue:
        layers.append([node.value for node in queue])
        queue = list(chain.from_iterable(node.subtrees for node in queue))

    return layers

def _get_layers_with_accumulator[A](
    queue: Collection[Tree[A]],
    acc: list[list[A]],
) -> list[list[A]]:
    if not queue:
        return acc

    return _get_layers_with_accumulator(
        tuple(chain.from_iterable(node.subtrees for node in queue)),
        [*acc, [node.value for node in queue]],
    )

def get_layers_recursive[A](tree: Tree[A]) -> list[list[A]]:
    """
    Perform a recursive breadth-first traversal, returning each layer as a
    sublist.
    """

    return _get_layers_with_accumulator([tree], [])

def get_layers_nonstrict[A](tree: Tree[A]) -> Generator[Generator[A]]:
    """
    Perform a breadth-first traversal, returning each layer as a sub-iterator.
    """

    queue = [tree]

    while queue:
        yield (node.value for node in queue)
        queue = list(chain.from_iterable(node.subtrees for node in queue))

def traverse_breadth[A](tree: Tree[A]) -> list[A]:
    """
    Perform a breadth-first traversal.
    """

    layers = []
    queue = [tree]

    while queue:
        layers.extend(node.value for node in queue)
        queue = list(chain.from_iterable(node.subtrees for node in queue))

    return layers

def _traverse_breadth_with_accumulator[A](
    queue: Collection[Tree[A]],
    acc: list[A],
) -> list[A]:
    if not queue:
        return acc

    return _traverse_breadth_with_accumulator(
        tuple(chain.from_iterable(node.subtrees for node in queue)),
        acc + [node.value for node in queue],
    )

def traverse_breadth_recursive[A](tree: Tree[A]) -> list[A]:
    """
    Perform a recursive breadth-first traversal.
    """

    return _traverse_breadth_with_accumulator([tree], [])

def traverse_breadth_nonstrict[A](tree: Tree[A]) -> Generator[A]:
    """
    Perform a breadth-first traversal, returning an iterator.
    """

    queue = [tree]

    while queue:
        yield from (node.value for node in queue)
        queue = list(chain.from_iterable(node.subtrees for node in queue))

class TestBfs(unittest.TestCase):
    TREE_SINGLE_ROOT = leaf(4)
    TREE_COMPLEX = Tree(3, [leaf(5), Tree(2, [leaf(6), leaf(1)]), leaf(7)])

    def test_get_layers_single_root(self):
        self.assertEqual(get_layers(self.TREE_SINGLE_ROOT), [[4]])
        
    def test_get_layers_complex(self):
        self.assertEqual(
            get_layers(self.TREE_COMPLEX),
            [[3], [5, 2, 7], [6, 1]],
        )

    def test_get_layers_recursive_single_root(self):
        self.assertEqual(get_layers_recursive(self.TREE_SINGLE_ROOT), [[4]])
        
    def test_get_layers_recursive_complex(self):
        self.assertEqual(
            get_layers_recursive(self.TREE_COMPLEX),
            [[3], [5, 2, 7], [6, 1]],
        )

    def test_get_layers_nonstrict_single_root(self):
        generator = get_layers_nonstrict(self.TREE_SINGLE_ROOT)

        self.assertEqual([list(layer) for layer in generator], [[4]])

    def test_get_layers_nonstrict_complex(self):
        generator = get_layers_nonstrict(self.TREE_COMPLEX)

        self.assertEqual(
            [list(layer) for layer in generator],
            [[3], [5, 2, 7], [6, 1]],
        )

    def test_traverse_breadth_single_root(self):
        self.assertEqual(traverse_breadth(self.TREE_SINGLE_ROOT), [4])

    def test_traverse_breadth_complex(self):
        self.assertEqual(
            traverse_breadth(self.TREE_COMPLEX),
            [3, 5, 2, 7, 6, 1],
        )

    def test_traverse_breadth_recursive_single_root(self):
        self.assertEqual(
            traverse_breadth_recursive(self.TREE_SINGLE_ROOT),
            [4],
        )

    def test_traverse_breadth_recursive_complex(self):
        self.assertEqual(
            traverse_breadth_recursive(self.TREE_COMPLEX),
            [3, 5, 2, 7, 6, 1],
        )

    def test_traverse_breadth_nonstrict_single_root(self):
        self.assertEqual(
            list(traverse_breadth_nonstrict(self.TREE_SINGLE_ROOT)),
            [4],
        )

    def test_traverse_breadth_nonstrict_complex(self):
        self.assertEqual(
            list(traverse_breadth_nonstrict(self.TREE_COMPLEX)),
            [3, 5, 2, 7, 6, 1],
        )

if __name__ == "__main__":
    unittest.main()