Question

I have a list where each element can be a number or one of two different strings (let's call 'a' and 'b'). I want to split the list in three lists: the list of numbers, the list of 'a' and the list of 'b'.

I can do it using a conditional list comprehension:

mixlist=[1,2,'a',3,'b','a']
numlist=[el for el in mixlist if type(el)==int]
alist=[el for el in mixlist if el=='a']
blist=[el for el in mixlist if el=='b']

but it leads to run several times on mixlist. Probably I can do better, but I do not know a short way to do it running over mixlist only one time...

Was it helpful?

Solution

You can do it straight forward, like this

numlist, alist, blist = [], [], []
for item in mixlist:
    if item == "a":
        alist.append(item)
    elif item == "b":
        blist.append(item)
    else:
        numlist.append(item)

Or you can shorten it a little bit, like this

numlist, alist, blist = [], [], []
d = {"a": alist, "b": blist}
for item in mixlist:
    d.get(item, numlist).append(item)

dict.get method will return the value corresponding to the key. But when the key is not found, it will return the default value passed as the second parameter. So, in our case, if the key is a or b we will get corresponding alist or blist. If neither of them, then we will get numlist and we append item to it.

Edit: You can slightly improve the straight forward way like this

mixlist = [1, 2, 'a', 3, 'b', 'a']

numlist, alist, blist = [], [], []
nappend, aappend, bappend = numlist.append, alist.append, blist.append

for item in mixlist:
    if item == "a":
        aappend(item)
    elif item == "b":
        bappend(item)
    else:
        nappend(item)
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top