Golang, Using Generator


  • 0
    Y

    How about this?

    func flatten(root *TreeNode)  {
    	if root == nil {
    		return
    	}
    	p := root
    	for node := range nodeGen(root) {
    		if node == p {
    			continue
    		}
    		p.Left = nil
    		p.Right = node
    		p = node
    	}
    }
    
    func nodeGen(root *TreeNode) <-chan *TreeNode {
    	c := make(chan *TreeNode)
    
    	go func() {
    		var walk func(r *TreeNode)
    		walk = func(r *TreeNode) {
    			if r == nil {
    				return
    			}
    			left, right := r.Left, r.Right
    			c <- r
    			walk(left)
    			walk(right)
    		}
    		walk(root)
    		close(c)
    	}()
    
    	return c
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.