» More on python flatten +

The standard (but for some unknown reason, not built-in) python flatten method looks something like this:

def flatten(l):
  out = []
  for item in l:
    if isinstance(item, (list, tuple)):
      out.extend(flatten(item))
    else:
      out.append(item)
  return out

This obviously runs into recursion errors pretty quickly for highly nested lists; what suprised me is that it can't grok even relatively shallow nesting, e.g., 10 levels deep:

a = []
for i in xrange(10):
  a = [a, i]
a = flatten(a)

Traceback (most recent call last):
  File "test.py", line 13, in ?
    a = flatten(a)
  File "test.py", line 5, in flatten
    out.extend(flatten(l))
  ...
  File "test.py", line 5, in flatten
    out.extend(flatten(l))
RuntimeError: maximum recursion depth exceeded

Someone smarter than I am (viz., Danny Yoo) wrote a better method using an iterator, which looks basically like this:

def iter_flatten(iterable):
  it = iter(iterable)
  for e in it:
    if isinstance(e, (list, tuple)):
      for f in iter_flatten(e):
        yield f
    else:
      yield e

a = []
for i in xrange(300):
  a = [a, i]
a = [i for i in iter_flatten(a)]

As you can see, this works with deeply nested arrays (up to 499 levels on my box). He also wrote a very opaque version using a form of tail recursion (via continuation passing), which is "not meant to be read by humans". That version can handle lists nested as deeply as the system recursion limit (1000 on my box)! That's pretty cool (even though my brain implodes when I try to understand it, heh)!

Just for fun, I decided to see what I could come up with. Here is my offering:

def flatten(l, limit=1000, counter=0):
  for i in xrange(len(l)):
    if (isinstance(l[i], (list, tuple)) and
        counter < limit):
      for a in l.pop(i):
        l.insert(i, a)
        i += 1
      counter += 1
      return flatten(l, limit, counter)
  return l

Nothing fancy. It's about as fast as Yoo's continuation version, but it breaks at 499 levels like the iterator version (but the iterator version is slower and requires the extra list-comprehension syntax if used for assignment). I also added an (optional) limit argument, to specify the maximum number of levels to flatten (ala ruby's Array#flatten). Don't worry about the counter argument, that's just to internally track state across recursions.

Addendum: Wow! I just came across a wonderful version of flatten in a cookbook comment. I saw it a few days ago, but I didn't really think about it, just kind of thought "yeah, another recursive flatten method," but this one whoops all of the others for speed, nesting level support and elegance! This is from Mike C. Fletcher's BasicTypes library. The method looks something like this (I've altered it a bit—see comments below):

def flatten(l, ltypes=(list, tuple)):
    ltype = type(l)
    l = list(l)
    i = 0
    while i < len(l):
        while isinstance(l[i], ltypes):
            if not l[i]:
                l.pop(i)
                i -= 1
                break
            else:
                l[i:i + 1] = l[i]
        i += 1
    return ltype(l)

a = []
for i in xrange(2000):
  a = [a, i]
a = flatten(a)

Freakin' genius!

* Fixed for empty lists/tuples based on Noah's comment
* Fixed again based on Greg's comment
* Fixed yet again based on John Y's comment

Addendum: I got to wondering how Mr. Fletcher's version would stack up against the built-in Array#flatten method in ruby. Granted, his doesn't have a flatten limit, and I think it would be kind of hard to add one, but then I have never really needed that feature. So here is the ruby version with times:

def flatten(l)
    i = 0
    while i < l.size
        while l[i].is_a? Array
            if l[i].empty?
                l.delete_at(i)
                i -= 1
                break
            else
                l[i...i+1] = l[i]
            end
        end
        i += 1
    end
    l
end

And now the times (tested with 1.8.6 final, best out of three). First, Fletcher's version:

p flatten(1500.times.inject { | m, i | m = [m, i] })

# time ruby test.rb
#
# real 0m0.038s
# user 0m0.028s
# sys 0m0.004s

Then the built-in version:

p 1500.times.inject { | m, i | m = [m, i] }.flatten

# time ruby test.rb
#
# real 0m0.027s
# user 0m0.016s
# sys 0m0.004s

Wow! That's pretty neat! A flatten implemented in ruby that is competitive with the C backend, heh! Now I know there are reasons for this, and like I said, Fletcher's version will just smash everything at every level, without regard. But still, it's always cool to find a bit of interpreted code that breaks out a can of Chuck Norris on the interpreter! (ruby is so manly it can almost beat itself up! heh!)