计算满三叉树的节点数

偶尔想到的一个问题，顺手写了一下解决方法。

具体实现

// 三叉树节点
type TripleTreeNode struct {
	Val    int
	Left   *TripleTreeNode
	Middle *TripleTreeNode
	Right  *TripleTreeNode
}

// 获取满三叉树的节点数
func getNodeCountOfFullTripleTree(root *TripleTreeNode) int64 {
	d := getDepthOfFullTripleTree(root)
	return getPowerOfN(3, int64(d)) / 2
}

// 获取满三叉树的深度。
// 空树深度是0。
// 只有根节点时，深度是1。
func getDepthOfFullTripleTree(root *TripleTreeNode) int {
	var d int
	curr := root
	for curr != nil {
		d++
		curr = curr.Left
	}
	return d
}

// 计算x的n次幂
// 输入要保证n>=1
func getPowerOfN(x, n int64) int64 {
	if n <= 0 {
		return 0
	}
	switch n {
	case 5:
		return x * x * x * x * x
	case 4:
		return x * x * x * x
	case 3:
		return x * x * x
	case 2:
		return x * x
	case 1:
		return x
	}
	if n&1 != 0 {
		y := getPowerOfN(x, n/2)
		return y * y * x
	}
	y := getPowerOfN(x, n/2)
	return y * y
}

测试代码

package main

import (
	"fmt"
	"math/rand"
	"time"
)

func main() {
	fmt.Println("Hello, playground")
	rand.Seed(time.Now().UnixNano())

	root := &TripleTreeNode{
		Val: 1,
	}
	maxDepth := 3
	AddChildren(root, maxDepth)
	fmt.Println(getNodeCountOfFullTripleTree(root))
}

func AddChildren(root *TripleTreeNode, depth int) {
	depth--
	if depth <= 0 {
		return
	}
	root.Left = &TripleTreeNode{
		Val: rand.Intn(1e5),
	}
	root.Middle = &TripleTreeNode{
		Val: rand.Intn(1e5),
	}
	root.Right = &TripleTreeNode{
		Val: rand.Intn(1e5),
	}
	AddChildren(root.Left, depth)
	AddChildren(root.Middle, depth)
	AddChildren(root.Right, depth)
}