Unival Trees in Go

unival.go

package main

//Given a tree, determine how many of its subtrees are inival trees.
//A unival tree is a tree where all nodes have the same value.
//An empty tree is a single unival tree.
//A tree with one node is a single unival tree.

//Assumptions
// - A tree whose value equals one of its children and the other child is nil is not a unival tree.
// - nil children do not contribute to the count of unival trees.

import "fmt"

type univalNode struct {
	value int
	left  *univalNode
	right *univalNode
}

func countUnivalSubTrees(root *univalNode) (isUnival bool, count int) {
	isUnival, count = true, 1
	if root == nil {
		return
	} else if root.left == nil && root.right == nil {
		return
	}
	leftIsUnival, leftCount := false, 0
	if root.left != nil {
		leftIsUnival, leftCount = countUnivalSubTrees(root.left)
	}
	rightIsUnival, rightCount := false, 0
	if root.right != nil {
		rightIsUnival, rightCount = countUnivalSubTrees(root.right)
	}
	isUnival = leftIsUnival && rightIsUnival && root.value == root.left.value && root.value == root.right.value
	count = leftCount + rightCount
	if isUnival {
		count++
	}
	return
}

func main() {
	fmt.Println(countUnivalSubTrees(nil))
	fmt.Println(countUnivalSubTrees(&univalNode{
		value: 0,
		left:  nil,
		right: nil,
	}))
	fmt.Println(countUnivalSubTrees(&univalNode{
		value: 0,
		left: &univalNode{
			value: 0,
			left:  nil,
			right: nil,
		},
		right: &univalNode{
			value: 1,
			left:  nil,
			right: nil,
		},
	}))
	fmt.Println(countUnivalSubTrees(&univalNode{
		value: 0,
		left: &univalNode{
			value: 0,
			left: &univalNode{
				value: 0,
				left: &univalNode{
					value: 0,
					left: &univalNode{
						value: 0,
						left:  nil,
						right: nil,
					},
					right: nil,
				},
				right: nil,
			},
			right: nil,
		},
		right: nil,
	}))
}