feat: add tree query algorithms (lowestCommonAncestor/distance/pathBetween/contains) (#35) (#42)

CHANGELOG.md

@@ -9,6 +9,8 @@ All notable changes to this project are documented here. The format is based on
 ### Added
 - Structural mutation helpers on `TreeNode`: `insertChild`, `removeChildAt`, `replaceChild`,
   `moveChild`, `addChildren`, and `sortChildren`.
+- Tree query extensions: `lowestCommonAncestor`, `distance`, `pathBetween`, and `contains` for
+  finding common ancestors, edge distances, the path between two nodes, and value membership.
 
 ### Changed
 - Rewrote the README for clarity: one consistent example tree, task-oriented sections

api/tree-structure.api

@@ -57,6 +57,13 @@ public final class com/github/adriankuta/datastructure/tree/TreeNodeNavigationEx
 	public static final fun siblings (Lcom/github/adriankuta/datastructure/tree/TreeNode;)Ljava/util/List;
 }
 
+public final class com/github/adriankuta/datastructure/tree/TreeNodeQueryExtKt {
+	public static final fun contains (Lcom/github/adriankuta/datastructure/tree/TreeNode;Ljava/lang/Object;)Z
+	public static final fun distance (Lcom/github/adriankuta/datastructure/tree/TreeNode;Lcom/github/adriankuta/datastructure/tree/TreeNode;)Ljava/lang/Integer;
+	public static final fun lowestCommonAncestor (Lcom/github/adriankuta/datastructure/tree/TreeNode;Lcom/github/adriankuta/datastructure/tree/TreeNode;)Lcom/github/adriankuta/datastructure/tree/TreeNode;
+	public static final fun pathBetween (Lcom/github/adriankuta/datastructure/tree/TreeNode;Lcom/github/adriankuta/datastructure/tree/TreeNode;)Ljava/util/List;
+}
+
 public final class com/github/adriankuta/datastructure/tree/TreeNodeSequenceExtKt {
 	public static final fun asSequence (Lcom/github/adriankuta/datastructure/tree/TreeNode;Lcom/github/adriankuta/datastructure/tree/iterators/TreeNodeIterators;)Lkotlin/sequences/Sequence;
 	public static synthetic fun asSequence$default (Lcom/github/adriankuta/datastructure/tree/TreeNode;Lcom/github/adriankuta/datastructure/tree/iterators/TreeNodeIterators;ILjava/lang/Object;)Lkotlin/sequences/Sequence;

src/commonMain/kotlin/com.github.adriankuta/datastructure/tree/TreeNodeQueryExt.kt

@@ -0,0 +1,88 @@
+package com.github.adriankuta.datastructure.tree
+
+/**
+ * The lowest (deepest) node that is an ancestor of both this node and [other], where every node is
+ * considered an ancestor of itself.
+ *
+ * Nodes are compared by identity (`===`), so this only returns a node when both arguments live in
+ * the same tree.
+ *
+ * @param other the other node to find the common ancestor with.
+ * @return the lowest common ancestor, or `null` when the two nodes belong to different trees and
+ *   therefore share no common ancestor.
+ *
+ * Runs in `O(da + db)` time and `O(da + db)` space, where `da`/`db` are the depths of the two nodes.
+ */
+public fun <T> TreeNode<T>.lowestCommonAncestor(other: TreeNode<T>): TreeNode<T>? {
+    // TreeNode has identity equality, so a HashSet gives O(1) identity membership and keeps the
+    // overall walk at O(da + db). Collect [other] and its ancestors, then climb from this node
+    // upward; the first node already on [other]'s chain is the deepest common ancestor.
+    val ancestorsOfOther = HashSet<TreeNode<T>>(other.ancestors())
+    ancestorsOfOther.add(other)
+    var node: TreeNode<T>? = this
+    while (node != null) {
+        if (node in ancestorsOfOther) return node
+        node = node.parent
+    }
+    return null
+}
+
+/**
+ * The number of edges on the shortest path between this node and [other].
+ *
+ * Computed as `depth() + other.depth() - 2 * lca.depth()`, where `lca` is their
+ * [lowestCommonAncestor]. The distance from a node to itself is `0`.
+ *
+ * @param other the other node to measure the distance to.
+ * @return the edge count, or `null` when the two nodes belong to different trees.
+ *
+ * Runs in `O(da + db)` time, where `da`/`db` are the depths of the two nodes.
+ */
+public fun <T> TreeNode<T>.distance(other: TreeNode<T>): Int? {
+    val lca = lowestCommonAncestor(other) ?: return null
+    return depth() + other.depth() - 2 * lca.depth()
+}
+
+/**
+ * The shortest path of nodes from this node to [other], inclusive of both endpoints.
+ *
+ * The path ascends from this node up to their [lowestCommonAncestor] and then descends to [other];
+ * the common ancestor appears exactly once. When `this === other` the result is `listOf(this)`. When
+ * one node is an ancestor of the other the path is simply the chain between them.
+ *
+ * @param other the node the path ends at.
+ * @return the path `[this, …, lca, …, other]`, or `null` when the two nodes belong to different
+ *   trees.
+ *
+ * Runs in `O(da + db)` time and space, where `da`/`db` are the depths of the two nodes.
+ */
+public fun <T> TreeNode<T>.pathBetween(other: TreeNode<T>): List<TreeNode<T>>? {
+    val lca = lowestCommonAncestor(other) ?: return null
+    val up = mutableListOf<TreeNode<T>>()
+    var node: TreeNode<T> = this
+    up.add(node)
+    while (node !== lca) {
+        node = node.parent!!
+        up.add(node)
+    }
+    val down = mutableListOf<TreeNode<T>>()
+    node = other
+    down.add(node)
+    while (node !== lca) {
+        node = node.parent!!
+        down.add(node)
+    }
+    return up + down.dropLast(1).reversed()
+}
+
+/**
+ * Returns `true` when this subtree contains a node whose value equals [value], including the
+ * receiver itself. Values are compared with `==` ([equals]).
+ *
+ * @param value the value to search for.
+ * @return `true` if any node in the pre-order traversal of this subtree holds [value].
+ *
+ * Runs in `O(n)` time over the `n` nodes of this subtree and stops at the first match.
+ */
+public fun <T> TreeNode<T>.contains(value: T): Boolean =
+    preOrderSequence().any { it.value == value }

src/commonTest/kotlin/com.github.adriankuta/datastructure.tree/TreeNodeQueryTest.kt

@@ -0,0 +1,123 @@
+package com.github.adriankuta.datastructure.tree
+
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertNull
+import kotlin.test.assertSame
+import kotlin.test.assertTrue
+
+class TreeNodeQueryTest {
+
+    // root(1)
+    // ├── n2(2)
+    // │   ├── n4(4)
+    // │   └── n5(5)
+    // └── n3(3)
+    //     └── n6(6)
+    private val root = TreeNode(1)
+    private val n2 = TreeNode(2)
+    private val n3 = TreeNode(3)
+    private val n4 = TreeNode(4)
+    private val n5 = TreeNode(5)
+    private val n6 = TreeNode(6)
+
+    // A completely separate tree.
+    private val otherRoot = TreeNode(10)
+    private val o11 = TreeNode(11)
+
+    init {
+        root.addChild(n2)
+        root.addChild(n3)
+        n2.addChild(n4)
+        n2.addChild(n5)
+        n3.addChild(n6)
+
+        otherRoot.addChild(o11)
+    }
+
+    @Test
+    fun lowestCommonAncestorOfTwoLeaves() {
+        assertSame(n2, n4.lowestCommonAncestor(n5))
+        assertSame(root, n4.lowestCommonAncestor(n6))
+    }
+
+    @Test
+    fun lowestCommonAncestorOfSameNode() {
+        assertSame(n4, n4.lowestCommonAncestor(n4))
+    }
+
+    @Test
+    fun lowestCommonAncestorOfAncestorAndDescendant() {
+        assertSame(n2, n2.lowestCommonAncestor(n4))
+        assertSame(n2, n4.lowestCommonAncestor(n2))
+        assertSame(root, root.lowestCommonAncestor(n6))
+    }
+
+    @Test
+    fun lowestCommonAncestorOfNodesInDifferentTreesIsNull() {
+        assertNull(n4.lowestCommonAncestor(o11))
+        assertNull(o11.lowestCommonAncestor(n4))
+    }
+
+    @Test
+    fun distanceValues() {
+        assertEquals(0, n4.distance(n4))
+        assertEquals(2, n4.distance(n5))
+        assertEquals(1, n2.distance(n4))
+        assertEquals(4, n4.distance(n6))
+        assertEquals(2, root.distance(n4))
+    }
+
+    @Test
+    fun distanceOfNodesInDifferentTreesIsNull() {
+        assertNull(n4.distance(o11))
+    }
+
+    @Test
+    fun pathBetweenSameNode() {
+        assertContentEquals(listOf(n4), n4.pathBetween(n4))
+    }
+
+    @Test
+    fun pathBetweenTwoLeaves() {
+        // n4 -> n2 -> n5 (lca = n2 appears once, endpoints are n4 and n5)
+        assertContentEquals(listOf(n4, n2, n5), n4.pathBetween(n5))
+        // n4 -> n2 -> root -> n3 -> n6 (lca = root appears once)
+        assertContentEquals(listOf(n4, n2, root, n3, n6), n4.pathBetween(n6))
+    }
+
+    @Test
+    fun pathBetweenWithUnequalDepthLegs() {
+        // Neither is an ancestor of the other and the legs differ in length: n4 is at depth 2, n3 at
+        // depth 1, lca = root. Exercises the asymmetric up/down assembly.
+        assertContentEquals(listOf(n4, n2, root, n3), n4.pathBetween(n3))
+        assertContentEquals(listOf(n3, root, n2, n4), n3.pathBetween(n4))
+    }
+
+    @Test
+    fun pathBetweenAncestorAndDescendant() {
+        assertContentEquals(listOf(n2, n4), n2.pathBetween(n4))
+        assertContentEquals(listOf(n4, n2), n4.pathBetween(n2))
+        assertContentEquals(listOf(root, n3, n6), root.pathBetween(n6))
+    }
+
+    @Test
+    fun pathBetweenOfNodesInDifferentTreesIsNull() {
+        assertNull(n4.pathBetween(o11))
+    }
+
+    @Test
+    fun containsTrueForValuesInSubtree() {
+        assertTrue(root.contains(1)) // the receiver itself
+        assertTrue(root.contains(6))
+        assertTrue(n2.contains(5))
+    }
+
+    @Test
+    fun containsFalseForValuesNotInSubtree() {
+        assertFalse(n2.contains(6)) // n6 lives under n3, not n2
+        assertFalse(root.contains(99))
+    }
+}