I am trying to implement a quick tool to find the combinations of a set of numbers (0...k), for j lots of this array, where the sum across the row is equal to k-j, and k>=j

For instance, for k=3, and j=2, I have all the following combinations:

[[0 0]
 [0 1]
 [0 2]
 [1 0]
 [1 1]
 [1 2]
 [2 0]
 [2 1]
 [2 2]]

There are only two elements that sum up to k-j=1:

[[0 1]
[1 0]]

My current implementation is slow as I approach high numbers of k and j:

        from itertools import product
        from numpy import 
        combs = np.array(list(product(np.arange(0, k), repeat=j)))
        combs = combs[np.sum(combs, axis=1)==k-j]
        print(combs)

Can anyone please suggest a more efficient algorithm than I have at the moment?

CodePudding user response：

For one, you don't have to consider the range up to k, but only up to k-j. Further, you could use itertools.combinations (if you are only interested in the set and not the order), as follows:

combs = np.array(list(combinations(range(k-j 1), j)))
combs = combs[np.sum(combs, axis=1)==k-j]