Generating evenly spaced triplets of probabilities in Python

I need to test different discrete probability distributions and it happens that I need triplets but it can probably be generalized to higher number of elements. I would like to cover the space quite evenly.

For example, for a step of 2, I would like something of the kind of :

 [[0,0,1],[0,0.5,0.5],[0,1,0],[0.5,0,0.5],[0.5,0.5,0],[1,0,0]]

My current idea is to generate all the triplets possibles so basically for a step 10 : [0,0,0] then [1/10, 0,0] ...

and reject the ones that do not add up to 1 but it is highly inefficient. Is there a better way to do so ?

I thought maybe some meshing technique or take a 1/n vector and then proceed to make sum of elements until they are 3 left.

Solution

You don't need to generate all triplets and then filter. Just generate 2 numbers, and subtract from 1.0 to get the third.

To expand on this idea:

import numpy as np


def get_triplets(n):
    a = np.linspace(0, 1, n + 1)
    a, b = np.meshgrid(a, a)
    c = 1 - a - b
    flattened = np.swapaxes(np.array([a, b, c]), 0, 2)
    # Check that each value of c is greater than zero
    flattened = flattened[flattened[:, :, 2] >= 0]
    return flattened


get_triplets(2)