Time: 208 ms (90.66%), Space: 38.4 MB (94.58%) - LeetHub

Author: hajin-kim

0501-find-mode-in-binary-search-tree/0501-find-mode-in-binary-search-tree.kt

@@ -1,20 +1,55 @@
 class Solution {
-    private val counter = mutableMapOf<Int, Int>()
-
-    private fun TreeNode?.count() {
-        if (this == null)
-            return
-        counter[this.`val`] = (counter[this.`val`] ?: 0) + 1
-        left.count()
-        right.count()
+    private fun morris(root: TreeNode): MutableList<Int> {
+        var currentStreak = 0
+        var currentVal = Int.MIN_VALUE
+        var maxStreak = 0
+        val result = mutableListOf<Int>()
+
+        var current: TreeNode? = root
+        while (current != null) {
+            if (current.left != null) {
+                // connect rightmost leaf of left subtree to current
+                // disconnect and move to left child
+
+                var prev = current.left!!
+                while (prev.right != null) {
+                    prev = prev.right!!
+                }
+
+                prev.right = current
+
+                val toVisit = current.left
+                current.left = null
+                current = toVisit
+            } else {
+                // count mode and traverse
+
+                if (current.`val` == currentVal) {
+                    ++currentStreak
+                } else {
+                    currentStreak = 1
+                    currentVal = current.`val`
+                }
+
+                if (currentStreak > maxStreak) {
+                    maxStreak = currentStreak
+                    result.clear()
+                    result.add(currentVal)
+                } else if (currentStreak == maxStreak) {
+                    result.add(currentVal)
+                }
+
+                current = current.right
+            }
+        }
+
+        return result
     }
 
     fun findMode(root: TreeNode?): IntArray {
-        root.count()
-        val max = counter.values.max()
-        return counter
-            .filter { (_, v) -> v == max }
-            .keys
-            .toIntArray()
+        if (root == null)
+            return IntArray(0)
+
+        return morris(root).toIntArray()
     }
 }