How can I count sequences that meet these constraints?

Question

I am trying to count permutations of a sequence of I and O symbols, representing e.g. people entering (I for "in") and leaving (O for "out") a room. For a given n many I symbols, there should be exactly as many O symbols, giving a total length of 2*n for the sequence. Also, at any point in a valid permutation, the number of O symbols must be less than or equal to the number of I symbols (since it is not possible for someone to leave the room when it is empty).

Additionally, I have some initial prefix of I and O symbols, representing people who previously entered or left the room. The output should only count sequences starting with that prefix.

For example, for n=1 and an initial state of '', the result should be 1 since the only valid sequence is IO; for n=3 and an initial state of II, the possible permutations are

IIIOOO
IIOIOO
IIOOIO

for a result of 3. (There are five ways for three people to enter and leave the room, but the other two involve the first person leaving immediately.)

I'm guessing the simplest way to solve this is using itertools.permutations. This is my code so far:

n=int(input())  ##actual length will be 2*n
string=input()
I_COUNT=string.count("I")
O_COUNT=string.count("O")
if string[0]!="I":
 sys.exit()
if O_COUNT>I_COUNT:
 sys.exit()
perms = [''.join(p) for p in permutations(string)]
print(perms)

the goal is to get the permutation for whatever is left out of the string and append it to the user's input, so how can I append user's input to the remaining length of the string and get the count for permutation?

Answer 1

The product() function from itertools will allow you to generate all the possible sequences of 'I' and 'O' for a given length.

From that list, you can filter by the sequences that start with the user-supplied start_seq.

From that list, you can filter by the sequences that are valid, given your rules of the number and order of the 'I's and 'O's:

from itertools import product


def is_valid(seq):
    '''Evaluates a sequence I's and O's following the rules that:
        - there cannot be more outs than ins
        - the ins and outs must be balanced
    '''
    _in, _out = 0, 0
    for x in seq:
        if x == 'I':
            _in += 1
        else:
            _out += 1

        if (_out > _in) or (_in > len(seq)/2):
            return False

    return True


# User inputs...
start_seq = 'II'

assert start_seq[0] != 'O', 'Starting sequence cannot start with an OUT.'

n = 3
total_len = n*2

assert len(start_seq) < total_len, 'Starting sequence is at least as big as total number, nothing to iterate.'

# Calculate all possible sequences that are total_len long, as tuples of 'I' and 'O'
seq_tuples = product('IO', repeat=total_len)

# Convert tuples to strings, e.g., `('I', 'O', 'I')` to `'IOI'`
sequences = [''.join(seq_tpl) for seq_tpl in seq_tuples]

# Filter for sequences that start correctly
sequences = [seq for seq in sequences if seq.startswith(start_seq)]

# Filter for valid sequences
sequences = [seq for seq in sequences if is_valid(seq)]

print(sequences)

and I get:

['IIIOOO', 'IIOIOO', 'IIOOIO']