0

I'm trying to write a code which "flattens" a list, it gets a list which contains several lists of numbers and creates one list of the numbers.

Running examples:

>>> flatten(1) 
[1] 
>>> flatten([]) 
[] 
>>> flatten([[[1],[2,[1]]]]) 
[1, 2, 1]

I want to do this recursively, but for some reason my code returns the same list:

def flatten(lst):
    if type(lst)==int:
        return [lst]
    elif lst==[]:
        return lst
    else:
        return [flatten(sub) for sub in lst]
msvalkon
  • 11,887
  • 2
  • 42
  • 38
Tam211
  • 723
  • 5
  • 10
  • 27

3 Answers3

0

There are two problems in your code

  1. instead of if type(lst)==int:, you should have written isinstance(lst, int)

  2. You are rebuilding an unwrapped list, so you actually have to chain the individual elements rather than wrapping in inside a list

    so, instead of return [flatten(sub) for sub in lst], it should be list(chain.from_iterable(flatten(sub) for sub in lst))

Your final code would look like

def flatten(lst):
    from itertools import chain
    if isinstance(lst,int):
        return [lst]
    elif lst==[]:
        return lst
    else:
        return list(chain.from_iterable(flatten(sub) for sub in lst))
Abhijit
  • 62,056
  • 18
  • 131
  • 204
  • @downvoter: Can you please leave a comment as to why you think this answer is not helpful? – Abhijit May 06 '14 at 10:55
  • IMHO `type(lst) is int` looks better than `isinstance(lst, int)`, but I agree that was a too minor reason for the downvote. Sorry, it was an impulsive decision of mine.Can you please edit your answer somehow, so I can retreat my downvote? – vaultah May 06 '14 at 11:29
  • @frostnational: the form `isinstance(obj,type)` or `isinstance(obj,(type1,..type_n))` is preferred to `type(obj) is type`. A) it might be a subclassed or virtual type you are testing, and your form will break; B) [isinstance](https://docs.python.org/2/library/functions.html#isinstance) supports a tuple of multiple types; C) `type(obj) is type` looks weird – dawg May 06 '14 at 14:21
0

Without using itertools

def flatten(items):
    for elem in items:
        if isinstance(elem,list):
            for sub_elem in flatten(elem):
                yield sub_elem
        else:
            yield elem

In [8]: list(flatten([[[1],[2,[1]]]]))
Out[8]: [1, 2, 1]
Padraic Cunningham
  • 176,452
  • 29
  • 245
  • 321
0

You could introduce result value with arguments as follows and retrieve the values recursively as follows:

def flatten(l, res=[]):
    if isinstance(l,list):
        for i in l:
            flatten(i,res)
        return res
    else:
        res.append(l)
        return res

print flatten([[[1],[2,[1]]]])

Output:

[1, 2, 1]
venpa
  • 4,268
  • 21
  • 23