Stacking numpy array A corresponding to array B

Question

I have a numpy array (ndarray, each of shape (1, 10, 4)) of input features X, and a label array y corresponding to classes.

Now I want to create a new array of arrays by stacking together, all input features in X that correspond to class 0 in y, features that correspond to class 1 in y, etc... so that a new array of arrays is formed, with nested the arrays equal to the number classes.

As a minimal example, say I have:

X = np.random.randn(200, 1, 10, 4)
a = np.zeros(100, dtype=int)
b = np.ones(100, dtype=int)
y = np.hstack((a,b))

So,

print(X.shape)
print(y.shape)
(200, 1, 10, 4)
(200,)

New array then should then be like:

Final = [array(#all_features_of class_0), array(#all_features_of_class_1)...]

My intention is to plot these features per class to understand their distribution.

If it may help, every single observation has 4 features, so the 4 in (1, 10, 4).

Answer 1

did I understood you correct, is this what you want?

class0, class1=[],[]
for i in range(len(y)):
if y[i]==0:
    class0.append(X[i])
else:
    class1.append(X[i])
Final = (np.array(class0) , np.array(class1))