From 49453b34b9505218ff427afee0f85fca5019abc1 Mon Sep 17 00:00:00 2001
From: Matej Focko <me@mfocko.xyz>
Date: Wed, 26 Jun 2024 14:33:37 +0200
Subject: [PATCH] =?UTF-8?q?go:=20add=20=C2=AB1382.=20Balance=20a=20Binary?=
 =?UTF-8?q?=20Search=20Tree=C2=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Matej Focko <me@mfocko.xyz>
---
 go/balance-a-binary-search-tree.go | 69 ++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)
 create mode 100644 go/balance-a-binary-search-tree.go

diff --git a/go/balance-a-binary-search-tree.go b/go/balance-a-binary-search-tree.go
new file mode 100644
index 0000000..1df8191
--- /dev/null
+++ b/go/balance-a-binary-search-tree.go
@@ -0,0 +1,69 @@
+package main
+
+func balanceBST(root *TreeNode) *TreeNode {
+	biggestPowerOfTwoThan := func(n int) int {
+		x := 1
+		for x <= n {
+			x <<= 1
+		}
+
+		return x >> 1
+	}
+
+	rotateLeft := func(parent, node *TreeNode) {
+		tmp := node.Right
+		node.Right = tmp.Left
+		tmp.Left = node
+		parent.Right = tmp
+	}
+	rotateRight := func(parent, node *TreeNode) {
+		tmp := node.Left
+		node.Left = tmp.Right
+		tmp.Right = node
+		parent.Right = tmp
+	}
+
+	rotate := func(head *TreeNode, count int) {
+		current := head
+		for i := 0; i < count; i++ {
+			tmp := current.Right
+			rotateLeft(current, tmp)
+			current = current.Right
+		}
+	}
+
+	if root == nil {
+		return nil
+	}
+
+	// Step 1 of DSW
+	head := TreeNode{Val: 0}
+	head.Right = root
+
+	current := &head
+	for current.Right != nil {
+		if current.Right.Left != nil {
+			rotateRight(current, current.Right)
+		} else {
+			current = current.Right
+		}
+	}
+
+	// Step 2 of DSW
+	count, current := 0, head.Right
+	for current != nil {
+		count++
+		current = current.Right
+	}
+
+	// Step 3 of DSW
+	m := biggestPowerOfTwoThan(count+1) - 1
+	rotate(&head, count-m)
+
+	for m > 1 {
+		m >>= 1
+		rotate(&head, m)
+	}
+
+	return head.Right
+}