Search code examples
pythonpython-2.7permutationpython-itertools

Python - Narrowing down itertools permutations results with "if" statements. How to with repeated element?


Relative python/programming newbie here. I am playing around with itertools permutations method to narrow down the yielded results with an "if" statement to create a list I need.

In my code below you will see that I want to return a narrowed down list from the permutations in which the indices of the commas match the indices of the commas of the list that is permutated. However, this code only works if I manually place the indices into the "if" statement. Since I will not know ahead of time how many commas will be in the list I want to permutate, how do I code for that?

from itertools import permutations

prp1 = ["UP", ",", "DOWN", "LEFT", ",", "RIGHT"]

# grab indices for commas in 'prp1' list
comma_indices = []
for index, p in enumerate(prp1):
    if p == ",":
        commas.append(index)

perms = permutations(prp1)
lst = []
for perm in perms:
    # Presently I have manually inserted the indices of commas,
    # but how to do this non-manual for unknown amount of commas here?
    if perm[1] == ',' and perm[4] == ',':
        lst.append(perm)

for l in lst:
    print l

I did think about doing something like this:

for perm in perms:
    for comma in commas:
        if...........

..but then I am only accessing 1 comma at a time.

Any help much appreciated, Darren

EDIT:

Some more example lists that will need to be permutated:

prp2 = ['down', ',', 'up', 'left', 'down']
prp3 = ['down', 'down', ',', 'up', ',', 'left', ',', 'right', ',', 'left']

Solution

  • You already know how to find indices where the commas are located without hardcoding:

    comma_indices = []
    for index, p in enumerate(prp1):
        if p == ",":
            commas.append(index)
    

    So you could apply the same approach to perm and see if it matches, e.g.

    observed_comma_indices = [index for index, p in enumerate(perm) if p == ","]
    if comma_indices == observed_comma_indices:
        lst.append(perm)
    

    But there are other approaches too. For example, instead of generating all the permutations and then keeping only the ones that you want, you could only create the ones you want in the first place:

    from itertools import permutations
    
    def permute_some(seq, indices_to_permute):
        for perm in permutations(indices_to_permute):
            newseq = seq[:]
            for i, p in zip(indices_to_permute, perm):
                newseq[i] = seq[p]
            yield newseq
    

    which gives

    >>> seq = ["A", "B", ",", "C", ","]
    >>> perm_ii = [i for i,x in enumerate(seq) if x != ","]
    >>> for p in permute_some(seq, perm_ii):
    ...     print p
    ...     
    ['A', 'B', ',', 'C', ',']
    ['A', 'C', ',', 'B', ',']
    ['B', 'A', ',', 'C', ',']
    ['B', 'C', ',', 'A', ',']
    ['C', 'A', ',', 'B', ',']
    ['C', 'B', ',', 'A', ',']