Generators

In my head, generators were a tool to make code faster. You had to think harder and write more complex code, and in exchange you got performance. More recently, I’ve come to realize that they can actually be easier to write than the equivalent eager code.

As part of testing out an interview question recently, I decided to write a function to generate k-length permutations of a given array. In Python, I might write something like

import itertools
itertools.permutations(range(10), 3)

but I was working in Javascript rather than Python and I wanted to implement the function myself. My Javascript was a bit rusty, with scoping rules and array methods that mutate the array while returning some other value both tripping me up repeatedly. After a few minutes of struggling to pass around my partial list of results, I decided to try generators. Using a generator, I was able to think about generating one permutation at a time, instead of generating a whole list of permutations at once. Eventually I came up with the following code.

function* generate(arr, k) {
  if (k == 0) {
    yield [];
  } else {
    for (const [i, el] of arr.entries()) {
      let remainder = arr.slice();
      remainder.splice(i, 1);
      for (let perm of generate(remainder, k-1)) {
        perm.unshift(el);
        yield perm;
      }
    }
  }
}

for (let perm of generate([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 3)) {
  console.log(perm);
}

I was excited enough about the result that I decided to write it in Python as well, to learn a bit about how generators in Python work.

def generate(arr, k):
  if k == 0:
    yield []
  else:
    for i, el in enumerate(arr):
      arr.pop(i)
      for perm in generate(arr, k-1):
        yield [el] + perm
      arr.insert(i, el)

print(list(generate(range(10), 3)))

First, if k == 0, we yield an empty list, since that is the only way to choose zero elements from a list. Note the else statement: unlike return, yield does not cause the function to exit, which can lead to some rather subtle bugs. For example, without the else condition (i.e. always descending into the for loop after the if statement), the function returns exactly the same results, but takes 12 seconds to run on my laptop instead of 0.04 seconds (with the arguments given above).

The code inside the for loop is relatively simple: we remove an element from the array to use as the first element of our permutation, then recursively generate the rest of the permutation and re-assemble it.

Once I had the lazy version of the code, writing the equivalent eager version became much easier.

def permutations(arr, k):
  if k == 0:
    return [[]]
  perms = []
  for i, el in enumerate(arr):
    arr.pop(i)
    for perm in permutations(arr, k-1):
      perms.append([el] + perm)
    arr.insert(i, el)
  return perms

print(permutations(range(10), 3))

Not having to think about building my full list of results and passing it around made writing the core recursion much easier, and then retrofitting the list- passing in was manageable. Even when I’m not using generators, I’ll look for opportunities to try generating one result, then generalizing that to produce a full list of results, especially when using recursion and when the core logic is non-trivial.