Combinations of M elements from N with non consecutive repetitions

Question

I have the following problem which can be summarized as follows:

Imagine you have two integers greater than 0, N (which defines the array n=np.array(range(N)) and M. We'd like to generate all the possible combinations of elements of n with length M, with the condition that no equal elements are consecutive.

For example, for N=3 (n=[0,1,2]) and M=3 we should obtain:

(0,1,0), (0,1,2) (0,2,0), (0,2,1), (1,0,1), (1,0,2), (1,2,0), (1,2,1), (2,0,1), (2,0,2), (2,1,0), (2,1,2)

That is, combinations such as (0,0,1), (1,1,1), (2,1,1)...etc, don't have to appear. Note that the number of all the valid combinations is simply given by N*(N-1)**(M-1).

As of now, for such example I am using this simple script (which also calculates all the combinations of varying length from m=1 to m=M):

import numpy as np
N = 3
M = 3
p = np.array(range(N))  

ic = [0]*M
c2 = np.zeros((int(N*(N-1)**(M-1)),M))
c1 = np.zeros((int(N*(N-1)**(M-2)),M-1))
c0 = np.zeros((int(N*(N-1)**(M-3)),M-2))
for i in p:
    c0[ic[0],:] = [i]
    ic[0] += 1
    for j in p[p!=i]:
        c1[ic[1],:] = [i,j]
        ic[1] += 1
        for k in p[p!=j]:
            c2[ic[2],:] = [i,j,k]
            ic[2] += 1

The problem is that this only works for this specific case with M=3, and M can be any integer greater than 0. So for some M, the previous code should have M nested loops which would have to be manually introduced.

I've tried defining a recursive function of variable number of loops, like this one that calculates the number of combinations (the number given by the equation above):

def rec_f(c,N,M):   
    if n>=1:
        for x in range(N):             
            c=rec_f(c,N,M-1)
    else:
        c += 1            
    return c

which I don't really even know why it works for that simple problem. Now, the thing is that I need to know the indices of the previous loops to be able replicate the script that generates all the possible combinations and I don't know how to that.

I've also tried to make one unique for loop (that would iterate N*(N-1)^(M-1) times), having in mind that the combinations can be expressed as numbers in base N, but after playing around for some time I got nothing useful.

I would appreciate any help, thanks in advance (and sorry for the long post)!

Answer 1

Just add the last element (if any) as an optional parameter to your recursive function. Also, no need for the N parameter, just pass the elements to chose from (also makes it more generally applicable). Also, I would recomment making it a generator function as the number of combinations might grow rather large, so you can consume them one by one as they come.

def combinations(elements, m, last=-1):
    if m:
        for x in elements:
            if x != last:
                for rest in combinations(elements, m-1, x):
                    yield (x,) + rest
    else:
        yield ()

Or a bit more compact, with yield from a generator expression:

def combinations(elements, m, last=-1):
    if m:
        yield from ((x,) + rest for x in elements if x != last
                                for rest in combinations(elements, m-1, x))
    else:
        yield ()

Example results, for both versions:

print(*combinations(range(3), 3))
# (0, 1, 0), (0, 1, 2), (0, 2, 0), (0, 2, 1), (1, 0, 1), (1, 0, 2), (1, 2, 0), (1, 2, 1), (2, 0, 1), (2, 0, 2), (2, 1, 0), (2, 1, 2)