user3483203
user3483203

Reputation: 51155

Split list into sub-lists based on attribute value

I have an array of objects that have a suit attribute, and I want to split into sub arrays based on which suit the object has. I currently am using this:

    for c in cards:
        if c.suit.value == 0:
            spades.append(c)
        elif c.suit.value == 1:
            diamonds.append(c)
        elif c.suit.value == 2:
            clubs.append(c)
        else:
            hearts.append(c)

I have tried to use itertools.groupby as follows:

suits = [list(g) for g in intertools.groupby(cards, lambda x: x.suit.value)]

But this just yields:

[[3, <itertools._grouper object at 0x000000000296B2E8>], ...]

My first approach works, I just imagine there is a simple pythonic one liner that accomplishes what I need.

Upvotes: 5

Views: 3560

Answers (3)

Jean-Fran&#231;ois Fabre
Jean-Fran&#231;ois Fabre

Reputation: 140168

When iterated upon, itertools.groupby yields tuple objects: the key, and an iterable of the grouped values (allowing lazy evaluation of the data).

You just converted the tuple to a list, which isn't very useful.

Instead, you need to drop the key (for instance by unpacking it to _, pythonic way of telling that you're not using it) and force iteration on the values:

[list(g) for _,g in itertools.groupby(cards, lambda x: x.suit.value)]

now if you want to group and sort, it's not very efficient to pass sorted(cards) to groupby just for the pleasure to create a one-liner (one-liners are nice, but I prefer fast programs). Your approach works, or you can also use collections.defaultdict, and even name the dict with the proper color using an indirection, for instance like this:

import collections

cards = collections.defaultdict(list)
colors = ["spades","diamonds","clubs","hearts"]

for c in cards:
    cards[colors[c.suit.value]].append(c)

Upvotes: 9

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476624

Although it is not a one-liner, by using a list, we make it more elegant:

spades, diamonds, clubs, hearts = collcard = [[] for _ in range(4)]

for c in cards:
    collcard[c.suit.value].append(c)

Here we thus initialize a list with four empty sublists, then we append the card c to the list with index c.suit.value.

We use iterable unpacking to assign the first element to spades, the second to diamonds, etc.

The advantage is that we avoid sorting (which works in O(n log n)). So this algorithm has time complexity O(n) (given the amortized cost of list appending is O(1)).

Although oneliners are usually elegant, one should not put to much effort in writing these, since oneliners can be harder to understand, or have significant impact with respect to performance.

Upvotes: 9

Netwave
Netwave

Reputation: 42698

If you know the size of the cards deck you can do something like this, considering is a 40 cards deck:

allcards = sorted(cards, key=lambda x: x.suit.value)
spades, diamonds, clubs, hearts =  map(lambda x: allcards[x:x+10], range(0,40,10))

Basiclly you get a sorted version of the deck, and take it in chuncks corresponding to the suits size. From spades to heards if they map from lower to higher index.

Upvotes: 1

Related Questions