yield vs. return - DFS binary tree problem

Question

Anyone please explain why 'return' isn't working here if I wrote if not node.left and not node.right: return node.val but working if the code goes:

class Solution:
def leafSimilar(self, root1: Optional[TreeNode], root2: Optional[TreeNode]) -> bool:
    def dfs(node):
        if not node:
            return None
        if not node.left and not node.right:
            yield node.val
        for i in dfs(node.left):
            yield i
        for i in dfs(node.right):
            yield i
    return all(a == b for a, b in itertools.zip_longest(dfs(root1), dfs(root2)))

with the example input root1 = [1,2,3], root2 = [1,3,2]

Why the result falsely shows 'True' in the 'return' case?

Answer 1

The function you've shown is a generator function, because it contains several yield statements. Even if you don't run any of those statements for certain inputs, their existence makes the whole function a generator function regardless.

When you call a generator function, the value you get back isn't its return value. Rather, it will be a generator object (a kind of iterator). You only get the return value when you consume the generator object with a yield from expression:

gen = some_generator_function()  # this assigns the generator object to gen
result = yield from gen          # this assigns the return value (after yielding values)

In the code you show, the calls to the generator function are being iterated on directly by for loops or indirectly by the itertools.izip_longest function. Nether of those approaches let you get a hold of the value returned by the generator function (though you probably could rewrite the loops to use yield from). Generally speaking, generator functions often don't return anything at all (which is equivalent to the return None that you're doing in the base case). It doesn't make much sense to return node.val in the current code. If you want to skip the loops (which will be over empty generators, which should be quick, but it's still doing a little useless work), you could yield node.val and then return on the next line (still within the if block).

Answer 2

In Python, the return statement is used to return a value from a function, and it immediately exits the function.

When you use return node.val, you are trying to return a value from the middle of a generator function. Generators use the yield statement to produce values one at a time and retain the state of the function between calls. Using return in this context doesn't make sense because it would exit the generator function immediately, preventing the subsequent values from being yielded.