华为od, 哈夫曼树

package main

import (
	"container/heap"
	"fmt"
)

type node struct {
	val   int
	child []*node
}

type IntHeap []*node

func (h IntHeap) Len() int { return len(h) }

func (h IntHeap) Less(i, j int) bool { return h[i].val < h[j].val }

func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x any) {
	*h = append(*h, x.(*node))
}

// Pop 要修改h的内容, 不用指针的花不会对原切片有改变
func (h *IntHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func main() {
	h := &IntHeap{&node{1, nil}, &node{2, nil}, &node{3, nil}}
	heap.Init(h)

	for h.Len() >= 2 {
		x := heap.Pop(h).(*node)
		y := heap.Pop(h).(*node)

		p := &node{val: x.val + y.val, child: []*node{x, y}}
		heap.Push(h, p)
	}

	root := heap.Pop(h).(*node)

	var dfs func(root *node)
	dfs = func(root *node) {
		if root.child == nil {
			fmt.Print(root.val, " ")
			return
		}

		dfs(root.child[0])
		fmt.Print(root.val, " ")
		dfs(root.child[1])
	}

	dfs(root)
}

全部评论

相关推荐

点赞 收藏 评论
分享
牛客网
牛客企业服务