上周D总在公司TTT上普及了go，当时只是概念上知道这门语言，D总出的小题目我用Java写了类似于DP的算法，被宣布太耗内存^^。这两天在公交上看了会reference（http://golang.org），最喜欢的特性是goroutines、多返回值和并列赋值/声明。觉得光看没用，还是写个helloworld吧，想到之前看到过一个这样的题目：
一颗二叉树，其节点上不均匀的分布了若干石头，石头数跟二叉树总节点数相同，石头只能在边上（即父子节点之间）进行搬运，每次只能搬运一颗石头。求使每个节点石头皆为一的最少搬运次数。
题目思路不多说了，甚至可能我的解法不是最优，这里主要是作为golang练手，在代码里详述。由于连机器上都还没装go环境，直接在http://play.golang.org上写的：
1 package main 2 3 import ( 4 "fmt" 5 ) 6 7 var N int = 9 8 var moves int = 0 9 10 type Node struct { 11 V int // number of stones the node keep 12 Feel int // number of stones the tree have, can be positive or negative 13 L *Node // left child 14 R *Node // right child 15 } 16 17 func (node *Node) String() string { 18 return fmt.Sprintf("{%d|%d}", node.V, node.Feel) 19 } 20 21 /** 22 what the mokced tree seems like: 23 0|1 24 / \ 25 1|3 2|1 26 / \ \ 27 3|0 4|1 5|0 28 \ / \ 29 6|1 7|0 8|2 30 */ 31 func mock() *Node { 32 // make new nodes 33 nodes := make([]*Node, N) 34 nodes[0] = &Node{1,0,nil,nil} 35 nodes[1] = &Node{3,0,nil,nil} 36 nodes[2] = &Node{1,0,nil,nil} 37 nodes[3] = &Node{0,0,nil,nil} 38 nodes[4] = &Node{1,0,nil,nil} 39 nodes[5] = &Node{0,0,nil,nil} 40 nodes[6] = &Node{1,0,nil,nil} 41 nodes[7] = &Node{0,0,nil,nil} 42 nodes[8] = &Node{2,0,nil,nil} 43 44 // construct tree 45 nodes[0].L, nodes[0].R = nodes[1], nodes[2] 46 nodes[1].L, nodes[1].R = nodes[3], nodes[4] 47 nodes[2].L, nodes[2].R = nil, nodes[5] 48 nodes[3].L, nodes[3].R = nil, nodes[6] 49 nodes[4].L, nodes[4].R = nil, nil 50 nodes[5].L, nodes[5].R = nodes[7], nodes[8] 51 nodes[6].L, nodes[6].R = nil, nil 52 nodes[7].L, nodes[7].R = nil, nil 53 nodes[8].L, nodes[8].R = nil, nil 54 return nodes[0] 55 } 56 57 /** 58 move stones between root, root.L, root.R, recursively. 59 */ 60 func move(root *Node) { 61 moves = 0 62 print("init", root) 63 count(root) 64 print("after count", root) 65 collect(root) 66 print("after collect", root) 67 welfare(root) 68 print("after welfare, finally got", root) 69 fmt.Println("all stone moves: ", moves) 70 } 71 72 /** 73 count feel of stones and number of nodes of root. 74 */ 75 func count(root *Node) (feel int) { 76 if root == nil { 77 return 0 78 } 79 feelL := count(root.L) 80 feelR := count(root.R) 81 root.Feel = feelL + feelR + root.V - 1 82 return root.Feel 83 } 84 85 /** 86 collect redundant stones up from rich child(ren) 87 */ 88 func collect(root *Node) { 89 if root == nil { 90 return 91 } 92 collect(root.L) 93 collect(root.R) 94 if root.L != nil && root.L.Feel > 0 { 95 // todo: number of stones to collect 96 todo := root.L.Feel 97 moves += todo 98 // move upward 99 root.V += todo 100 root.L.Feel = 0 101 root.L.V -= todo 102 } 103 if root.R != nil && root.R.Feel > 0 { 104 todo := root.R.Feel 105 moves += todo 106 root.V += todo 107 root.R.Feel = 0 108 root.R.V -= todo 109 } 110 } 111 112 /** 113 dispatch all stones collected to poor child(ren) 114 */ 115 func welfare(root *Node) { 116 if root == nil { 117 return 118 } 119 if root.L != nil && root.L.Feel < 0 { 120 todo := -root.L.Feel 121 root.L.Feel = 0 122 root.L.V += todo 123 root.Feel = 0 124 root.V -= todo 125 moves += todo 126 } 127 if root.R != nil && root.R.Feel < 0 { 128 todo := -root.R.Feel 129 root.R.Feel = 0 130 root.R.V += todo 131 root.Feel = 0 132 root.V -= todo 133 moves += todo 134 } 135 welfare(root.L) 136 welfare(root.R) 137 } 138 139 /** 140 bfs print using chan as queue 141 */ 142 func print(title string, root *Node) { 143 fmt.Println("==========| ", title, " |=========="); 144 queue := make(chan *Node, N) 145 queue <- root 146 i := 0 147 for node := range queue { 148 fmt.Print(node, "; ") 149 if node.L != nil { 150 queue <- node.L 151 } 152 if node.R != nil { 153 queue <- node.R 154 } 155 if i += 1; i == N { 156 close(queue) 157 break 158 } 159 } 160 fmt.Println() 161 } 162 163 func main() { 164 move(mock()) 165 }
其中最自我感觉良好的是，自己想到了个主意用chan作为queue把二叉树BFS打印出来^^，虽然这个在熟手们可能是常用做法。
代码里略写了些注释。
如果：
- 思路/逻辑不正确；
- 还有比这更优化的解法；
- 有些代码不符合go规范或约定；
请不吝点评：D
有疑问加站长微信联系（非本文作者）
