diff options
| -rw-r--r-- | spec/binary_trees/kth_largest_in_bst_spec.rb | 10 |
1 files changed, 6 insertions, 4 deletions
diff --git a/spec/binary_trees/kth_largest_in_bst_spec.rb b/spec/binary_trees/kth_largest_in_bst_spec.rb index 15aa99e..65449cb 100644 --- a/spec/binary_trees/kth_largest_in_bst_spec.rb +++ b/spec/binary_trees/kth_largest_in_bst_spec.rb @@ -153,16 +153,18 @@ describe "#kth_largest_in_bst" do Traversal.new(k).traverse(tree) end - def dfs(node, yielder) + def traverse(node, yielder) return if node.nil? - dfs(node.left, yielder) + traverse(node.left, yielder) yielder.yield node.value - dfs(node.right, yielder) + traverse(node.right, yielder) end def kth_largest_in_bst(tree, k) - Enumerator.new { |yielder| dfs(tree, yielder) }.take(k).last + x = Enumerator.new { |yielder| traverse(tree, yielder) } + (k-1).times { x.next } + x.next end [ |