Find Leaves of Binary Tree

Question

Working on following problem:

Given a binary tree, collect a tree's nodes as if you were doing this: Collect and remove all leaves, repeat until the tree is empty.

Example:
Given binary tree 
          1
         / \
        2   3
       / \     
      4   5    
Returns [4, 5, 3], [2], [1].

Explanation:
1. Removing the leaves [4, 5, 3] would result in this tree:

          1
         / 
        2          
2. Now removing the leaf [2] would result in this tree:

          1          
3. Now removing the leaf [1] would result in the empty tree:

          []         
Returns [4, 5, 3], [2], [1].

My idea was a simple recursive algorithm shown below. The idea is to find the leaves of the left subtree and the right subtree, and weave them such that the depths are in the right subarray. I've tested the 'weave' method pretty thoroughly, and I think it's fine. My concern is with my recursive implementation-- I'm getting an answer way off from the correct one, and not sure why.

Below is my code with sample input/output:

def find_leaves(root)
    return [] if root.nil?
    #create leaf_arr of root.left and root.right
    #weave them in order.
    #add the root

    left_arr = find_leaves(root.left)
    right_arr = find_leaves(root.right)


    weave(left_arr, right_arr) << [root]
end


def weave(arr1, arr2) #these are 2d arrs
    i = 0
    until i == arr1.length || i == arr2.length #potential nil/empty case here
        arr1[i] += arr2[i]
        i += 1
    end
    if i < arr2.length  
        #either arr 1 or arr2 isn't finished. if arr1 isn't finished, we're done. if arr2 isnt finished, do the below:
        until i == arr2.length
            arr1 << arr2[i]
            i += 1
        end
    end
    arr1
end

Sample input/output/correct answer:

Run Code Result: ×

input: [1,2,3,4,5]

Your answer: [[[4],[5],[3]],[[2,4,5]],[[1,2,3,4,5]]]

Expected answer: [[4,5,3],[2],[1]]

I've printed the output for the left_arr and right_arr variables and they look fine, and I've stress-tested my weave algorithm. Am I off conceptually here?

Answer 1

Since this still sits open and seems to fair highly when I google search your title. I'll show a pretty expressive solution:

def find_leaves(root)
  return [] if root.nil?
  return [[root.val]] if root.left.nil? && root.right.nil?
  todo = [root]
  leaves = []

  until todo.empty?
    top = todo.shift
    %w[left right].each do |path|
      leaf = top.send(path)
      next if leaf.nil?
      if leaf.left.nil? && leaf.right.nil?
        leaves << leaf.val
        top.instance_variable_set("@#{path}", nil)
      else
        todo << leaf
      end
    end
  end
  [leaves].concat(find_leaves(root))
end

A more refactored version:

def find_leaves(root)
  leaves = []
  search = lambda do |branch|
    return -1 unless branch
    i = 1 + [search[branch.left], search[branch.right]].max
    (leaves[i] ||= []) << branch.val
    i
  end
  search[root]
  leaves
end

They're both about the same speed, and really the first one is easier to read and understand.

Answer 2

In your code you are using pure depth first search algorithm DFS and with that algorithm I think that you can hardly achieve your goal with array joggling you are doing in weave function. Because your tree will be processed in this order 4 , 5 , 2 , 3 , 1. One solution will be to do it with iteration (pseudo code):

function doJob(root) begin
  leaves = findLeaves(root)
  while leaves.size > 0 do begin
    for each leaf in leaves delete(leaf)
    leaves = findLeaves(root)
  end
  delete(root)
end

function findLeaves(node) begin
  if node = nil then begin
    return []
  end
  else begin
    leftLeaves = findLeaves(node.left)
    rightLeaves = fingLeaves(node.right)
    leaves = leftLeaves + rightLeaves
    if leaves.size == 0 then begin
      leaves.add(node)
    end
    return leaves
  end
 end

Answer 3

I can't comment so I will do it like this. (do remember that i dont know ruby) I think something goes already wrong in how the double arrays (root.left and root.right) are defined. How are they defined? how is root defined?

But the following eplains the repeat of the whole array.

weave(left_arr, right_arr) << [root]   

This should be someting in the line of this.

weave(left_arr, right_arr) << [root.root]

Otherwise you are appending the whole root array wich is [1,2,3,4,5]. So this explains the adding of last part. [[[4],[5],[3]],[[2,4,5]],[[1,2,3,4,5]]].

My suggestion in finding the error in weave would be to print arr1 and arr2 at every stage.... Could you show that..