CODEWITHSUNDEEP
pythonlist-comprehension
Geoff
Geoff

Reputation: 1007

How to include a count condition within a list comprehension?

I have created the following:

from itertools import product
x = [(b0, b1, b2, b3) for b0, b1, b2, b3 in product(range(5), repeat=4)]

This creates all tuples from [0,0,0,0] to [4,4,4,4].

I would like as a condition to only include those tuples which do not have the same number repeated more than 2 times. So I would like to ignore such tuples as [2,2,2,1] or [3,3,3,3] while keeping tuples such as [0,0,1,2] or [1,3,4,2]

I tried the following and I think it is ok but it appears rather cumbersome. 

y = [(b0, b1, b2, b3) for b0, b1, b2, b3 in product(range(5), repeat=4) if (b0, b1, b2, b3).count(0)<=2 and (b0, b1, b2, b3).count(1)<=2 and (b0, b1, b2, b3).count(2)<=2 and (b0, b1, b2, b3).count(3)<=2  and (b0, b1, b2, b3).count(4)<=2]

Perhaps a way to count each of the elements from [0,1,2,3,4] and take their max and state that the max<=2.

How is it possible to include with a list comprehension a counting condition?

Upvotes: 1

Views: 78

Answers (2)

JohanL
JohanL

Reputation: 6891

Using set might work. Another option is to use collections.Counter:

from collections import Counter
from itertools import product
x =  [
     comb for comb in product(range(5), repeat=4)
     if max(Counter(comb).values()) <= 2
     ]

Upvotes: 3

Thierry Lathuille
Thierry Lathuille

Reputation: 24232

You could create a generator to produce your tuples, and check your condition with a Counter:

from itertools import product
from collections import Counter

def selected_tuples():
    for t in product(range(5), repeat=4):
        if Counter(t).most_common(1)[0][1]<=2:
            yield t

print(list(selected_tuples()))

Output:

[(0, 0, 1, 1), (0, 0, 1, 2), (0, 0, 1, 3), (0, 0, 1, 4), ...]

Upvotes: 2

Related Questions