Criticize my first attempt at using the itertools module

Question

Here's the exercise:

Given a string of X's and O's and a maximum number of swaps, determine the longest possible sequence of X's after swaps.

Here's my solution using itertools.

import itertools as iter

def swaps_gen(s):
    '''This generator yields all the possible outcome of applying one swap
    to the input string'''
    X_positions = [i[0] for i in enumerate(s) if i[1] == "X"]
    O_positions = [i[0] for i in enumerate(s) if i[1] == "O"]

    for x in X_positions:
        tmp_x = X_positions[:]
        tmp_o = O_positions[:]

        tmp_x.remove(x)
        tmp_o.append(x)

        for o in O_positions:
            ttmp_x = tmp_x[:]
            ttmp_o = tmp_o[:]

            ttmp_o.remove(o)
            ttmp_x.append(o)

            s = []
            for i in ttmp_x: s.insert(i, "X")
            for i in ttmp_o: s.insert(i, "O")
            yield "".join(s)

def flatten(l):
    '''flatten one level. [[a], [b, c]] --> [a, b, c]'''
    return iter.chain.from_iterable(l) 

def max_x_after_swap(s, n):
    '''repeat the generate swap_gen n times.
    I used a set() to avoid duplicates.
    Then returns the highest number of consecutive "X"s it finds.'''
    outcomes = [s]
    while n > 0:
        outcomes = set(flatten([list(swaps_gen(x)) for x in outcomes]))
        n -= 1
    return max(map(nbr_of_consecutive_x, outcomes))

def nbr_of_consecutive_x(s):
    '''return the number of consecutive "X" '''
    return max([len(list(v)) for k, v in iter.groupby(s) if k == "X"])

def main():
    input = "XXOXOXOXO"
    swaps_nbr = 2
    print(max_x_after_swap(input, swaps_nbr))

main()

Any criticism is welcome. It'd be great if you could point out any anti-pattern I may present, or show me better (by better I mean more pythonic) way to do things. I am more interested in improving my Python style than improving my algorithm.

I used Python3.1 to run the above.

Answer 1

import itertools as iter

I dislike abbrevations like this. Itertools is widely used within the python community, and it only hinders code readability by using iter instead of itertools

def swaps_gen(s):
    '''This generator yields all the possible outcome of applying one swap
    to the input string'''
    X_positions = [i[0] for i in enumerate(s) if i[1] == "X"]
    O_positions = [i[0] for i in enumerate(s) if i[1] == "O"]

PEP 8 recommends lowercase_with_underscores for local variable names. Your comprehensions would be more clearly written as

X_positions = [index for index, value in enumerate(s) if value == 'X']

Because you do it twice, I'd consider writing a function

    for x in X_positions:
        tmp_x = X_positions[:]
        tmp_o = O_positions[:]

I despise uses of tmp in variable names. All variables are temporary.

        tmp_x.remove(x)
        tmp_o.append(x)

That's going to be expensive. The remove has to move the whole list back to make it fit.

        for o in O_positions:
            ttmp_x = tmp_x[:]
            ttmp_o = tmp_o[:]

Especially egregious is using tmp again with an extra t.

            ttmp_o.remove(o)
            ttmp_x.append(o)

            s = []

I wouldn't reuse s from before

            for i in ttmp_x: s.insert(i, "X")
            for i in ttmp_o: s.insert(i, "O")

Inserts may be expensive. It won't be so bad here because your inserts will be in mostly sorted order. But the one unsorted element in each list will cost you somewhat.

            yield "".join(s)

Your whole approach here is a bit awkward. I'd suggest something like

for x_index in x_positions:
    for y_index in y_positions:
         swapped_list = list(s) # copy string into list, so I can modify it
         swapped_list[x_index], swapped_list[y_index] = swapped_list[y_index],swapped_list[x_index] # swap elements
         yield ''.join(swapped_list) # back to list

Shorter and probably more efficient.

def flatten(l):
    '''flatten one level. [[a], [b, c]] --> [a, b, c]'''
    return iter.chain.from_iterable(l) 

def max_x_after_swap(s, n):
    '''repeat the generate swap_gen n times.
    I used a set() to avoid duplicates.
    Then returns the highest number of consecutive "X"s it finds.'''

PEP 257 recommends using """ not ''' for docstrings. Also, the closing ''' should be on a line by itself

    outcomes = [s]

I'd make this a set, just for consistency with the next line

    while n > 0:

Why a while loop instead of a for loop?

        outcomes = set(flatten([list(swaps_gen(x)) for x in outcomes]))

I'd use outcomes = set(flatten(map(swaps_gen, outcomes))) There is no reason to use list here, as the generator will flatten just fine.

        n -= 1
    return max(map(nbr_of_consecutive_x, outcomes))

def nbr_of_consecutive_x(s):

I dislike abbreviation like nbr, I'd call this count_of_consecutive_x

    '''return the number of consecutive "X" '''
    return max([len(list(v)) for k, v in iter.groupby(s) if k == "X"])

You don't need the [], because max will take a generator expression just fine

def main():
    input = "XXOXOXOXO"
    swaps_nbr = 2
    print(max_x_after_swap(input, swaps_nbr))

main()