python - numpy 3d array -- flatten --> 1d array --> select one element in 1d --> how to know the index of the element in 3d?

Question

numpy 3d array -- flatten --> 1d array --> select one element in 1d array --> how to know the index of the element in 3d? (?, ?, ?)

For example,

I have 3d histogram data in 3d-array D. Let D be D = np.random.randint(0,100, size=(4,3,2)).

and then the probability of each element would be equal to P = D / D.sum(0).sum(0).sum(0).

Let P be flattened like P_flat = P.reshape(-1).

and then randomly chose n-sample from P_flat: samples = np.random.choice(len(P_flat), n, p=P_flat)

and then, I hope to convert 1D index (=samples) to 3D index of D.

For example, convert(samples = [3, 2, 5]) -> [[?,?,?], [?,?,?], [?,?,?].

Answer 1

You can use np.unravel_index (doc) for that:

indices = np.array(np.unravel_index(samples, D.shape)).T