Kotlin binary tree breadth first - fully lazy traversal

lazy-breadth-first-trversal.kt

/* related to: https://x.com/vikhyatk/status/1873033432705712304 */

import java.util.LinkedList

data class BinaryTreeNode(val label: String, val left: BinaryTreeNode? = null, val right: BinaryTreeNode? = null)

data class NodeAtLevel(val level: Int, val node: BinaryTreeNode)

/**
 *  Lazy Breadth First Traversal of the binary tree by node, left to right by level
 */
fun lazyBreadthFirstTraversal(rootNode: BinaryTreeNode): Sequence<NodeAtLevel> {
    val queue = LinkedList<NodeAtLevel>()
    fun BinaryTreeNode.queueAt(level: Int) {
        queue += NodeAtLevel(level, this)
    }

    rootNode.queueAt(1)
    return generateSequence {
        queue.removeFirstOrNull()?.also { current ->
            current.node.left?.queueAt(current.level + 1)
            current.node.right?.queueAt(current.level + 1)
        }
    }
}

data class LazyNodesAtLevel(val level: Int, val nodes: Sequence<BinaryTreeNode>)

/**
 * Full lazy traversal of the tree grouped by level (a sequence of levels, each a sequence of nodes)
 */
fun lazyBreadthFirstTraversalByLevel(rootNode: BinaryTreeNode): Sequence<LazyNodesAtLevel> = sequence {
    val traversal = lazyBreadthFirstTraversal(rootNode).withLookahead().iterator()
    if (!traversal.hasNext()) return@sequence
    var lastLevel = 0

    while (traversal.hasNext()) {
        val (first, next) = traversal.next()
        if (first.level == lastLevel) continue; // the caller did not drain the inner sequence, so we will
        lastLevel = first.level;

        val levelSequence = LazyNodesAtLevel(first.level, sequence {
            yield(first.node)
            while (next != null && next.level == first.level && traversal.hasNext()) {
                val (node, nextNode) = traversal.next()
                yield(node.node)
                if (nextNode?.level != first.level) break
            }
        })
        yield(levelSequence)
    }
}

data class NodesAtLevel(val level: Int, val nodes: List<BinaryTreeNode>)

/**
 *  Partially lazy traversal of the tree grouped and materialized by level (a sequence of levels, each a list of nodes)
 */
fun breadthFirstTraversalByLevel(rootNode: BinaryTreeNode): Sequence<NodesAtLevel> =
    lazyBreadthFirstTraversalByLevel(rootNode)
        .map { (level, nodes) -> NodesAtLevel(level, nodes.toList()) }

private fun <T> Sequence<T>.withLookahead(): Sequence<Pair<T,T?>> = this.windowed(2,1,true).map { it.get(0) to it.getOrNull(1) }

fun main() {
    val tree = BinaryTreeNode("One",
        BinaryTreeNode("two", BinaryTreeNode("four")),
        BinaryTreeNode("three", BinaryTreeNode("five"), BinaryTreeNode("six"))
    )

    println("Lazy breadth First Traversal by Node:")
    lazyBreadthFirstTraversal(tree).forEach { println(it) }
    println()

    println("Lazy Breadth First Traversal by Level:")
    lazyBreadthFirstTraversalByLevel(tree).forEach {
        println("${it.level} => ${it.nodes.joinToString { it.label }}")
    }
    println()

    println("Lazy Breadth First Traversal by Level, skipping level 2:")
    lazyBreadthFirstTraversalByLevel(tree).filter { it.level != 2 }.forEach {
        println("${it.level} => ${it.nodes.joinToString { it.label }}")
    }
    println()

    println("Lazy Breadth First Traversal by Level, first node only per level:")
    lazyBreadthFirstTraversalByLevel(tree).forEach {
        println("${it.level} => ${it.nodes.first().label}")
    }
    println()


    println("Eager Breadth First Traversal by Level (lazy by Node):")
    breadthFirstTraversalByLevel(tree).forEach {
        println("${it.level} => ${it.nodes.joinToString { it.label }}")
    }
    println()
}