From c662de0f8317ad343f5dba3ce17996c9fd44af3b Mon Sep 17 00:00:00 2001
From: motomux <motomux@gmail.com>
Date: Fri, 24 Feb 2017 23:59:31 -0800
Subject: [PATCH] Add problem 3 of chapter 4

---
 src/chapter4/binary_tree_node.go |  21 +++++
 src/chapter4/problem3.go         |  40 +++++++++
 src/chapter4/problem3_test.go    | 145 +++++++++++++++++++++++++++++++
 3 files changed, 206 insertions(+)
 create mode 100644 src/chapter4/binary_tree_node.go
 create mode 100644 src/chapter4/problem3.go
 create mode 100644 src/chapter4/problem3_test.go

diff --git a/src/chapter4/binary_tree_node.go b/src/chapter4/binary_tree_node.go
new file mode 100644
index 0000000..9751ad8
--- /dev/null
+++ b/src/chapter4/binary_tree_node.go
@@ -0,0 +1,21 @@
+package chapter4
+
+// BTNode represents a node in binary tree
+type BTNode struct {
+	Value  int
+	Parent *BTNode
+	Left   *BTNode
+	Right  *BTNode
+}
+
+// SetRight sets child on right side
+func (n *BTNode) SetRight(c *BTNode) {
+	n.Right = c
+	c.Parent = n
+}
+
+// SetLeft sets child on right side
+func (n *BTNode) SetLeft(c *BTNode) {
+	n.Left = c
+	c.Parent = n
+}
diff --git a/src/chapter4/problem3.go b/src/chapter4/problem3.go
new file mode 100644
index 0000000..400f5dd
--- /dev/null
+++ b/src/chapter4/problem3.go
@@ -0,0 +1,40 @@
+package chapter4
+
+import "container/list"
+
+// ListOfDepth returns slice of list which has nodes on each depth
+func ListOfDepth(root *BTNode) []*list.List {
+	if root == nil {
+		return nil
+	}
+
+	var res []*list.List
+	queue := list.New()
+	queue.PushFront(root)
+	for queue.Len() > 0 {
+		res = append(res, cloneList(queue))
+		nextQueue := list.New()
+		for queue.Len() > 0 {
+			elm := queue.Back()
+			queue.Remove(elm)
+			btNode := elm.Value.(*BTNode)
+
+			if btNode.Left != nil {
+				nextQueue.PushFront(btNode.Left)
+			}
+			if btNode.Right != nil {
+				nextQueue.PushFront(btNode.Right)
+			}
+		}
+
+		queue = nextQueue
+	}
+
+	return res
+}
+
+func cloneList(l *list.List) *list.List {
+	clone := list.New()
+	clone.PushBackList(l)
+	return clone
+}
diff --git a/src/chapter4/problem3_test.go b/src/chapter4/problem3_test.go
new file mode 100644
index 0000000..d57ff67
--- /dev/null
+++ b/src/chapter4/problem3_test.go
@@ -0,0 +1,145 @@
+package chapter4
+
+import (
+	"container/list"
+	"reflect"
+	"testing"
+)
+
+func TestListOfDepth(t *testing.T) {
+	tests := map[string]struct {
+		in  *BTNode
+		out []*list.List
+	}{
+
+		"Should return slice of a single list when binary tree has only root": {
+			in: &BTNode{
+				Value: 1,
+			},
+			out: []*list.List{
+				pushList(list.New(), &BTNode{Value: 1}),
+			},
+		},
+
+		"Should return slice of lists when binary tree has nodes": {
+			in: &BTNode{
+				Value: 4,
+				Left: &BTNode{Value: 2,
+					Left:  &BTNode{Value: 1},
+					Right: &BTNode{Value: 3},
+				},
+				Right: &BTNode{Value: 6,
+					Left:  &BTNode{Value: 5},
+					Right: &BTNode{Value: 7},
+				},
+			},
+			out: []*list.List{
+				pushList(list.New(),
+					&BTNode{
+						Value: 4,
+						Left: &BTNode{Value: 2,
+							Left:  &BTNode{Value: 1},
+							Right: &BTNode{Value: 3},
+						},
+						Right: &BTNode{Value: 6,
+							Left:  &BTNode{Value: 5},
+							Right: &BTNode{Value: 7},
+						},
+					}),
+				pushList(list.New(),
+					&BTNode{Value: 2,
+						Left:  &BTNode{Value: 1},
+						Right: &BTNode{Value: 3},
+					},
+					&BTNode{Value: 6,
+						Left:  &BTNode{Value: 5},
+						Right: &BTNode{Value: 7},
+					},
+				),
+				pushList(list.New(),
+					&BTNode{Value: 1},
+					&BTNode{Value: 3},
+					&BTNode{Value: 5},
+					&BTNode{Value: 7},
+				),
+			},
+		},
+
+		"Should return slice of lists when binary tree is unbalanced": {
+			in: &BTNode{
+				Value: 4,
+				Left:  &BTNode{Value: 2},
+				Right: &BTNode{Value: 6,
+					Left: &BTNode{Value: 5},
+					Right: &BTNode{Value: 7,
+						Right: &BTNode{Value: 8,
+							Right: &BTNode{Value: 9},
+						},
+					},
+				},
+			},
+			out: []*list.List{
+				pushList(list.New(),
+					&BTNode{
+						Value: 4,
+						Left:  &BTNode{Value: 2},
+						Right: &BTNode{Value: 6,
+							Left: &BTNode{Value: 5},
+							Right: &BTNode{Value: 7,
+								Right: &BTNode{Value: 8,
+									Right: &BTNode{Value: 9},
+								},
+							},
+						},
+					}),
+				pushList(list.New(),
+					&BTNode{Value: 2},
+					&BTNode{Value: 6,
+						Left: &BTNode{Value: 5},
+						Right: &BTNode{Value: 7,
+							Right: &BTNode{Value: 8,
+								Right: &BTNode{Value: 9},
+							},
+						},
+					},
+				),
+				pushList(list.New(),
+					&BTNode{Value: 5},
+					&BTNode{Value: 7,
+						Right: &BTNode{Value: 8,
+							Right: &BTNode{Value: 9},
+						},
+					},
+				),
+				pushList(list.New(),
+					&BTNode{Value: 8,
+						Right: &BTNode{Value: 9},
+					},
+				),
+				pushList(list.New(), &BTNode{Value: 9}),
+			},
+		},
+
+		"Should return nil if root is nil": {
+			in:  nil,
+			out: []*list.List(nil),
+		},
+	}
+
+	for k, test := range tests {
+		t.Run(k, func(t *testing.T) {
+			out := ListOfDepth(test.in)
+			if !reflect.DeepEqual(out, test.out) {
+				t.Errorf("actual=%+v expected=%+v", out, test.out)
+			}
+		})
+	}
+}
+
+func pushList(l *list.List, data ...interface{}) *list.List {
+	for _, d := range data {
+		l.PushFront(d)
+	}
+
+	return l
+}