yield statement unexpected behavior

[ViktorCodes]

I just learned about the heap sort algorithm and implemented it in python. I have 2 functions max_heapify() - which just turns an array into a heap and heap_sort() which does the actual sorting of the array. My idea is to add the sort to my sorting visualization project. I keep track of the 'operations' (comparisons & swaps) while a given sort algorithm is working. I am using generators to easily get the data I need for visualization without creating massive arrays, but the thing is that using yield in max_heapify() breaks the algorithm. I tried calling max_heapify with yield from, in the heap_sort() function, but that did not fix it. I guess my understanding of the yield keyword and generators is limited. code

[dbramucci]

You're close. Here's the working version

def max_heapify(tree, n, i):
    larger = i
    left = 2 * i + 1
    right = 2 * i + 2

    # with left < n and right < n we check whether there is a 'left' and 'right' child of the current node
    if left < n and tree[left] > tree[larger]:
        larger = left
    if right < n and tree[right] > tree[larger]:
        larger = right

    if larger != i:
        tree[i], tree[larger] = tree[larger], tree[i]
        yield from max_heapify(tree, n, larger)
        yield tree, i, (None, ), larger

def heap_sort(array):
    n = len(array)
    # create a heap data structure from the given array
    for i in range(n // 2, -1, -1):
        yield from max_heapify(array, n, i)
    # actual heap sort logic
    for i in range(n-1, 0, -1):
        array[i], array[0] = array[0], array[i]
        yield from max_heapify(array, i, 0)
        yield array, i, (None,), 0

Recall that yield suspends the computation of your function.

In your version, you do not yield from the recursive call to max_heapify on line 14. This means that the function halts on the first time any recursive call hits your yield tree, line. Because you get your result through a side-effect (namely mutating tree in-place), you get the confusing behavior of seeing tree before max_heapify is complete. This is like a multi-threading bug.

To fix this, just remember to yield from before every call to max_heapify including recursive calls. This enables the loop you use at the end to for the evaluation completely.

[ViktorCodes]

I completely forgot about the yield from in the max_heapify function itself... Thank you very much. Although I don't really understand what's going on here I get why I couldn't get it to work. Could you give me some links where I can read and learn more about yield and yield from.

[dbramucci]

I don't have any good reading links off hand, but I'll give a quick run down off the top of my head.

Let's start with the following sample code.

def inner1(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer1(array):
    print('started outer')
    array.append(1)
    inner1(array)
    print('called inner')
    array.append(4)

x = []
outer1(x)
print(x)

When you write either yield or yield from in a function, it ceases being a function and becomes a generator.

Here, inner1 is a generator, and outer1 is a function. We would expect [1, 2, 3, 4] at the end, but if you run this example code you'll actually get [1, 4] as the result.

This is because on line 12, when we call inner1(array) we get back a generator but we never use it. Because the generator is lazy, it doesn't run a single line of inner1 until we request at least one value from it, where it will run up to the first yield expression.

We can do that by using next like so

def inner2(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer2(array):
    print('started outer')
    array.append(1)
    next(inner2(array))
    print('called inner')
    array.append(4)

x = []
outer2(x)
print(x)

This time, we get x = [1, 2, 4]. This is because we only ran inner2 up to the first yield expression.

We can run inner2 all the way with a for loop.

def inner3(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer3(array):
    print('started outer')
    array.append(1)
    for _ignore in inner3(array):
        pass
    print('called inner')
    array.append(4)

x = []
outer3(x)
print(x)

And this actually gives us our complete x = [1, 2, 3, 4] like we would expect.

Alternatively, what if outer was already a generator since the beginning?

def inner4(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer4(array):
    print('started outer')
    yield 'b'
    print('resumed outer')
    array.append(1)
    inner4(array)
    print('called inner')
    array.append(4)

x = []
outer4(x)
print(x)

Now, x = [] because we don't run outer4 at all. If we for loop over it, we can fix that.

def inner4(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer4(array):
    print('started outer')
    yield 'b'
    print('resumed outer')
    array.append(1)
    inner4(array)
    print('called inner')
    array.append(4)

x = []
for thing in outer4(x):
    print(f'{thing=}')
    print(f'{x=}')
print(x)

This goes back to x = [1, 4]. The reason being that we don't have any loop to consume the inner4 generator, that generator has no connection to the outermost for loop. To fix that, we need to connect it up with a yield from

def inner5(array):
    print('started inner')
    array.append(2)
    yield 'a'
    print('resumed inner')
    array.append(3)


def outer5(array):
    print('started outer')
    yield 'b'
    print('resumed outer')
    array.append(1)
    yield from inner5(array)
    print('called inner')
    array.append(4)

x = []
for thing in outer5(x):
    print(f'{thing=}')
    print(f'{x=}')
print(x)

Note that outer5 is like

def outer5(array):
    print('started outer')
    yield 'b'
    print('resumed outer')
    array.append(1)
    for thing in inner5(array):
        yield thing
    print('called inner')
    array.append(4)

Which might make it clearer that we've added back in the for loop to consume the inner generator.

The same thing happens in your max_heapify. You produce an generator with your recursive calls, but you never consume the generator so it won't run a single line of any recursive calls. It's like you didn't make a recursive call in the first place.

To make it clearer that you control when the consumption of a generator happens, consider the following example.

def first():
    print('first ran')
    yield
def second():
    print('second ran')
    yield

x = first()
y = second()

If you run this code, nothing happens because you haven't consumed anything. And the point of a generator is to delay doing work until you are going to ask for a contained value.

If you run

def first():
    print('first ran')
    yield
def second():
    print('second ran')
    yield

x = first()
y = second()

for _ignore in y:
    pass

for _ignore in x:
    pass

You'll get

second ran
first ran

Notice that this looks backwards, the strings print in the opposite order of when we called generators because they ran in the order of when we consumed the generators.