Python Itertools Code Optimisation

Question

Given question (the contest is now over)

• a password consists of exactly n lowercase English letters.
• the password is melodious, meaning that consonants can only be next to vowels and vowels can only be next to consonants. Example: bawahaha
• the password cannot contain the letter y (because it's both a consonant and vowel).
• the first letter of the password can be either a vowel or consonant.

Given the length, n, of the password,

print all of the possible passwords that meet the conditions above.

Input Format

The line of input contains the integer (the length of the password).

Constraints

Output Format

Print each of the possible passwords, one per line. The order of the passwords does not matter.

My Code in Python:

import sys
import itertools

n = int(raw_input().strip())

consonants = ['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't', 'v', 'w', 'x', 'z']
vowels = ['a', 'e', 'i', 'o', 'u']


test4 = ['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't', 'v', 'w', 'x', 'z', 'a', 'e', 'i', 'o', 'u']
answer = set(itertools.product(test4, repeat=n))

for letters in answer:
    for j in xrange(len(letters)):
        flag = 1
        if j != len(letters) - 1:
            if letters[j] in vowels and letters[j+1] in vowels:
                flag = 0
                break
            if letters[j] in consonants and letters[j+1] in consonants:
                flag = 0
                break

    if flag:
        for j in letters:
            sys.stdout.write(j)
        print ""

Is there a better way to do this?

Answer 1

Of course there's a better way (if you mean faster). You can generate a itertools.product where you don't have to "discard" items.

You can simply create a product of vowel, consonant, vowel, consonant, ... (alternating both lists n times) and one which starts with consonant, vowel, consonant, vowel, .... The returned items will always satisfy the condition so all that needs to be done is printing them.

import itertools

def alternating(seq1, seq2, length):
    """Generator that alternatingly yields seq1 and seq2, `length` times"""
    assert length >= 0
    for _ in range(length // 2):
        yield seq1
        yield seq2
    if length % 2 == 1:
        yield seq1

consonants = ['b', 'c', 'd', 'f', 'g', 'h', 'j', 'k', 'l', 'm', 'n', 'p', 'q', 'r', 's', 't', 'v', 'w', 'x', 'z']
vowels = ['a', 'e', 'i', 'o', 'u']

n = int(raw_input().strip())

# Starts with vowel
for comb in itertools.product(*list(alternating(vowels, consonants, n))):
    print(''.join(comb))

# Starts with consonant
for comb in itertools.product(*list(alternating(consonants, vowels, n))):
    print(''.join(comb))

That way you can reduce the number of possible candidates.

Your approach gave 25**n items, while the new approach only generates 2 * 5**(n//2)*20**(n//2) (if n even) or 5**(n//2 + 1) * 20 ** (n//2) + 5**(n//2) * 20 ** (n//2 + 1) (if n odd) items.

For n=5 that means: What generated originally 9765625 items (almost 10 million!) from product will now only generate 250000 items. Even ignoring the (possibly very expensive) check if your sequence satisfies the problem-condition (which is obsolete now) you generated 40 times more items in your product!