Solution for "Flatten Binary Tree to Linked List"

Author: Thuy Trinh

build.gradle

@@ -1,5 +1,5 @@
 buildscript {
-  ext.kotlinVersion = '1.2.0'
+  ext.kotlinVersion = '1.2.10'
 
   repositories {
     mavenCentral()

src/main/java/leetcode/binarytreeboundary/TreeNode.java

@@ -1,8 +1,12 @@
 package leetcode.binarytreeboundary;
 
+import org.jetbrains.annotations.Nullable;
+
 public class TreeNode {
   public int val;
+  @Nullable
   public TreeNode left;
+  @Nullable
   public TreeNode right;
 
   public TreeNode(int x) {

src/main/java/leetcode/flattenbtree/Solution.kt

@@ -0,0 +1,26 @@
+package leetcode.flattenbtree
+
+import leetcode.binarytreeboundary.TreeNode
+
+/**
+ * https://leetcode.com/problems/flatten-binary-tree-to-linked-list/description/
+ */
+class Solution {
+  fun flatten(root: TreeNode?) {
+    root?.flatten()
+  }
+
+  private fun TreeNode.flatten(): TreeNode? {
+    if (left == null && right == null) {
+      return this
+    }
+
+    val newLeftLeaf = left?.flatten()
+    val newRightLeaf = right?.flatten()
+    val oldRight = right
+    right = left ?: oldRight
+    left = null
+    newLeftLeaf?.right = oldRight
+    return newRightLeaf ?: newLeftLeaf ?: this
+  }
+}

src/test/java/leetcode/flattenbtree/SolutionTest.kt

@@ -0,0 +1,44 @@
+package leetcode.flattenbtree
+
+import leetcode.binarytreeboundary.isEqualTo
+import leetcode.binarytreeboundary.node
+import org.amshove.kluent.shouldBeTrue
+import org.junit.Test
+
+class SolutionTest {
+  @Test
+  fun case0() {
+    val root = node(1) {
+      left = node(2) {
+        left = node(3)
+        right = node(4)
+      }
+      right = node(5) {
+        right = node(6)
+      }
+    }
+    Solution().flatten(root)
+    root.isEqualTo(node(1) {
+      right = node(2) {
+        right = node(3) {
+          right = node(4) {
+            right = node(5) {
+              right = node(6)
+            }
+          }
+        }
+      }
+    }).shouldBeTrue()
+  }
+
+  @Test
+  fun case1() {
+    val root = node(1) {
+      left = node(2)
+    }
+    Solution().flatten(root)
+    root.isEqualTo(node(1) {
+      right = node(2)
+    }).shouldBeTrue()
+  }
+}