Generating all possible combinations of n-sized vector that follow certain conditions on each element

Question

I have a list d of length r such that d = (d_1, d_2,..., d_r). I would like to generate all possible vectors of length r such that for any i (from 0 to r), v_i is between 0 and d_i.

For example,

if r =2 and d= (1,2), v_1 can be 0 or 1 and v_2 can be 0,1 or 2. 

Hence there are 6 possible vectors: [0,0] , [0,1], [0,2], [1,0] , [1,1], [1,2]

I have looked into Itertools and combinations and I have a feeling I will have to use recursion however I have not managed to solve it yet and was hoping for some help or advice into the right direction.

Edit: I have written the following code for my problem and it works however I did it in a very inefficient way by disregarding the condition and generating all possible vectors then pruning the invalid ones. I took the largest d_i and generated all vectors of size r from (0,0,...0) all the way to (max_d_i,max_d_i,....max_d_i) and then eliminated those that were invalid.

Code:

import itertools
import copy
def main(d):
    arr = []
    correct_list =[]
    curr = []
    r= len(d)
    greatest = max(d)
    for i in range(0,greatest+1):
        arr = arr + [i]
    #all_poss_arr is a list that holds all possible vectors of length r from (0,0,...,0) to (max,max,...,max)
    # for example if greatest was 3 and r= 4, all_poss_arr would have (0,0,0,0), then (0,0,0,1) and so on,
    #all the way to (3,3,3,3)
    all_poss_arr = list(itertools.product(arr,repeat = r))    
    #Now I am going to remove all the vectors that dont follow the v_i is between 0 and d_i
    for i in range(0,len(all_poss_arr)):
        curr = all_poss_arr[i]
        cnt = 0
        for j in range(0,len(curr)):
            if curr[j] <= d[j]:
                cnt = cnt +1
        if cnt == r:
            curr = list(curr)
            currcopy = copy.copy(curr)
            correct_list = correct_list + [currcopy]
            cnt =0
    return correct_list

If anyone knows a better way, let me know, it is much appreciated.

Answer 1

You basically want a Cartesian product. I'll demonstrate a basic, functional and iterative approach.

Given

import operator as op
import functools as ft
import itertools as it


def compose(f, g):
    """Return a function composed of two functions."""
    def h(*args, **kwargs):
        return f(g(*args, **kwargs))
    return h


d = (1, 2)

Code

Option 1: Basic - Manual Unpacking

list(it.product(range(d[0] + 1), range(d[1] + 1)))
# [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

Option 2: Functional - Automated Mapping

def vector_combs(v):
    """Return a Cartesian product of unpacked elements from `v`."""
    plus_one = ft.partial(op.add, 1)
    range_plus_one = compose(range, plus_one)
    res = list(it.product(*map(range_plus_one, v)))
    return res


vector_combs(d)
# [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

Option 3: Iterative - Range Replication (Recommended)

list(it.product(*[range(x + 1) for x in d]))
# [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

---

Details

Option 1

The basic idea is illustrated in Option 1:

• Make a Cartesian product using a series of modified ranges.

Note, each range is manually incremented and passed in as an index from d. We automate these limitations in with the last options.

Option 2

We apply a functional approach to handle the various arguments and functions:

• Partial the 1 argument to the add() function. This returns a function that will increment any number.
• Let's pass this function into range through composition. This allows us to have a modified range function that auto increments the integer passed in.
• Finally we map the latter function to each element in tuple d. Now d works with any length r.

Example (d = (1, 2, 1), r = 3):

vector_combs((1, 2, 1))
# [(0, 0, 0),
#  (0, 0, 1),
#  (0, 1, 0),
#  (0, 1, 1),
#  (0, 2, 0),
#  (0, 2, 1),
#  (1, 0, 0),
#  (1, 0, 1),
#  (1, 1, 0),
#  (1, 1, 1),
#  (1, 2, 0),
#  (1, 2, 1)]

Option 3

Perhaps most elegantly, just use a list comprehension to create r ranges. ;)