CODEWITHSUNDEEP
python
marshall
marshall

Reputation: 2483

Iterate over lists with a certain number of 1s, 0s and -1s

I can iterate over all lists with -1s, 0s and 1s with

for v in itertools.product([-1,0,1], repeat = n):

However, if I only want lists with A 1s, B 0s and C -1s, how can I iterate over those without making all the lists and filtering with if v.count(1)=A and v.count(0) = B and v.count(C)=-1 ?

EDIT

Using itertools.permutations is very wasteful sadly as it makes the same tuple over and over again.

As len(list(itertools.permutations([1]*2 + [0]*2 + [-1]*2))) = 720 although len(set(itertools.permutations([1]*2 + [0]*2 + [-1]*2))) = 90.

And we can see that permutations repeats tuples with

 print list(itertools.permutations([1]*1 + [0]*1 + [-1]*2))
[(1, 0, -1, -1), (1, 0, -1, -1), (1, -1, 0, -1), (1, -1, -1, 0), (1, -1, 0, -1), (1, -1, -1, 0), (0, 1, -1, -1), (0, 1, -1, -1), (0, -1, 1, -1), (0, -1, -1, 1), (0, -1, 1, -1), (0, -1, -1, 1), (-1, 1, 0, -1), (-1, 1, -1, 0), (-1, 0, 1, -1), (-1, 0, -1, 1), (-1, -1, 1, 0), (-1, -1, 0, 1), (-1, 1, 0, -1), (-1, 1, -1, 0), (-1, 0, 1, -1), (-1, 0, -1, 1), (-1, -1, 1, 0), (-1, -1, 0, 1)]

Upvotes: 0

Views: 154

Answers (2)

Paul Hankin
Paul Hankin

Reputation: 58339

Using itertools.permutations will produce duplicates. You can code it up yourself, and here's one way to do it.

def uniq_perms(a, b, c):
    if a < 0 or b < 0 or c < 0:
        return
    if a + b + c == 0:
        yield []
    for s in uniq_perms(a - 1, b, c):
        yield [0] + s
    for s in uniq_perms(a, b - 1, c):
        yield [1] + s
    for s in uniq_perms(a, b, c - 1):
        yield [-1] + s

for s in uniq_perms(2, 1, 1):
    print s

Upvotes: 3

Inbar Rose
Inbar Rose

Reputation: 43477

What you want is actually permutations not product.

Create a function to wrap this functionality:

def foo(A, B, C):
    return itertools.permutations([1]*A + [0]*B + [-1]*C)

Usage Example:

>>> for v in foo(1,1,1):
    print v

(1, 0, -1)    (-1, 1, 0)
(-1, 0, 1)    (0, 1, -1)
(0, -1, 1)    (1, -1, 0)


>>> for v in foo(2,1,1):
    print v

(1, 1, 0, -1)    (1, 0, 1, -1)    (0, -1, 1, 1)
(1, 1, -1, 0)    (1, 0, -1, 1)    (0, -1, 1, 1)
(1, 0, 1, -1)    (1, -1, 1, 0)    (-1, 1, 1, 0)
(1, 0, -1, 1)    (1, -1, 0, 1)    (-1, 1, 0, 1)
(1, -1, 1, 0)    (0, 1, 1, -1)    (-1, 1, 1, 0)
(1, -1, 0, 1)    (0, 1, -1, 1)    (-1, 1, 0, 1)
(1, 1, 0, -1)    (0, 1, 1, -1)    (-1, 0, 1, 1)
(1, 1, -1, 0)    (0, 1, -1, 1)    (-1, 0, 1, 1)

Explanation:

What you want to do is create a list which will be used to produce all the permutations, and you want this list to contain an amount of -1,0,1s as you want. Lets start with creating a list full of 5 1's. We can do that with:

>>> [1]*5
[1, 1, 1, 1, 1]

We can add this list to another list:

>>> [1]*5 + [0]*2
[1, 1, 1, 1, 1, 0, 0]

And yet another:

>>> [1]*5 + [0]*2 + [-1]*7
[1, 1, 1, 1, 1, 0, 0, -1, -1, -1, -1, -1, -1, -1]

And so, when you want to create your list, we get:

[1]*A + [0]*B + [-1]*C

Upvotes: 1

Related Questions