Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

elegantly splitting a list (or dict) into two via some arbitrary function in python

Is there any elegant way of splitting a list/dict into two lists/dicts in python, taking in some arbitrary splitter function?

I could easily have two list comprehensions, or two selects, but it seems to me there should be some better way of doing it that avoids iterating over every element twice.

I could do it easily with a for loop and if statement, but that takes something like 7 lines of code for what should be a very simple operation.

Any ideas?

Edit:

Just for reference, my two solutions would be,

# given dict cows, mapping cow names to weight
# fast solution
fatcows = {}
thincows = {}
for name, weight in cows:
    if weight < 100:
        thincows[name] = weight
    else:
        fatcows[name] = weight

# double-list-comprehension solution would be
fatcows = {name: weight for name, weight in cows.items() if weight > 100}
thincows = {name: weight for name, weight in cows.items() if weight < 100}

I was thinking there must be something more elegant than this that i never thought of, something like:

thincows, fatcows = ??? short expression involving cows ???

I know it's possible to do by writing higher order functions stuff to do it for me, and i know how to do it manually. I was just wondering if there was some super-elegant language feature to do it for me.

It's like you can write your own subroutines and whatnot to do a SELECT on a list, or you can just say

thincows = select(cows, lambda c: c.weight < 100)

I was hoping there would be some similarly elegant way of splitting the list, with one pass

like image 250
Li Haoyi Avatar asked Jul 31 '11 02:07

Li Haoyi


2 Answers

How about 3 lines?

fatcows, thincows = {}, {}
for name, weight in cows.items():
    (fatcows if weight > 50 else thincows)[name] = weight

Tested:

>>> cows = {'bessie':53, 'maud':22, 'annabel': 77, 'myrna':43 }
>>> fatcows, thincows = {}, {}
>>> for name, weight in cows.items():
...     (fatcows if weight > 50 else thincows)[name] = weight
... 
>>> fatcows
{'annabel': 77, 'bessie': 53}
>>> thincows
{'maud': 22, 'myrna': 43}
like image 104
senderle Avatar answered Oct 06 '22 00:10

senderle


Any solution is going to take O(N) time to compute, whether it be through two passes through the list or one pass that does more work per item. The simplest way is just to use the tools that are available to you: itertools.ifilter and itertools.ifilterfalse:

def bifurcate(predicate, iterable):
    """Returns a tuple of two lists, the first of which contains all of the
       elements x of `iterable' for which predicate(x) is True, and the second
       of which contains all of the elements x of `iterable` for which
       predicate(x) is False."""
    return (itertools.ifilter(predicate, iterable),
            itertools.ifilterfalse(predicate, iterable))
like image 22
Adam Rosenfield Avatar answered Oct 05 '22 23:10

Adam Rosenfield