查找树的某一层

以前有朋友问我怎么查找树的某一层节点, 还花了点时间去思考, 觉得不应该这种基本问题还不能一下子反应过来, 于是手动用go实现了一遍. 直接看代码:

package main

import "fmt"
import "container/list"

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

func find_level_nodes(root *Node, level int) (ret []int) {
    if root == nil || level < 1 {
        return nil
    }

    queue := list.New()
    queue.PushBack(root)
    cur_level := 0
    next_first_node := root

    for queue.Len() > 0 {
        f := queue.Front()
        node := f.Value.(*Node)

        if node == next_first_node {
            cur_level++
            if cur_level > level {
                break
            }

            next_first_node = nil
        }

        if node.left != nil {
            queue.PushBack(node.left)
            if next_first_node == nil {
                next_first_node = node.left
            }
        }
        if node.right != nil {
            queue.PushBack(node.right)
            if next_first_node == nil {
                next_first_node = node.right
            }
        }

        fmt.Println(node.value)
        if cur_level == level {
            ret = append(ret, node.value)
        }

        queue.Remove(f)
    }

    return ret
}

func main() {
    fmt.Println("hello")

    node1 := &Node{
        value: 1,
    }

    /*
        node2 := &Node{
            value: 2,
        }
    */

    node3 := &Node{
        value: 3,
        left:  node1,
    }

    node4 := &Node{
        value: 4,
    }

    node1.right = node4

    ret := find_level_nodes(node3, 2)
    fmt.Println(ret)
}

主要思想还是在于用队列来实现层遍历, 以及标记下一层初始节点来计算层次, 其中下一层的初始节点未必能首次找到, 只要在遍历当前层次节点的时候发现下一层初始节点为空就可以设置.