yの値に準じてxを分類したデータセットを作成したい

Question

前提・実現したいこと

numpy配列のデータAX,　AYがあります。

AX（200列×100行）
[0.11, 0.12, 0.33....
1.42, 1.43. 1.65...
0.52, 0.24, 0.735...
1.68, 1.39, 1.49, 1....
...]
AY(100列)
[0, 1, 0, 1,...]

AYの値(0または1の値を取る)によってAXを新しいリストA1,A0
の２つに分類したいです。AXの行とAYの列はデータとして対応しています。

#理想
A1　#AYが1の時のAX
[1.42, 1.43. 1.65...
1.68, 1.39, 1.49,...
...]
A0 #AYが0の時のAX
[0.11, 0.12, 0.33....
0.52, 0.25, 0.735...
...]

発生している問題・エラーメッセージ

ValueError: operands could not be broadcast together with shapes (0,) (200,) 

また、実際にこのforループが正しく機能しているのかも確認できません。

該当のソースコード

A1 = A0 = []
for i in range(len(AX)):
      if AY[i] is 1:
        A1 += AX[i]
      else:
        A0 += AX[i]
A1_test, A0_test = A1, A0
...

初歩的なことで申し訳ございません。よろしくお願いいたします。

Answer 1

例えばですが、こんな感じでどうでしょうか。numpyのwhereを使ってインデックスを取り出して、該当する行列を取り出せば簡単だと思います。

対象のデータ(AX, AY)は自前でランダムに生成しました。

import numpy as np

AX = np.random.rand(100, 200)
AY = np.random.randint(100)

row_sum = 2
col_sum = 1
AY     = np.zeros((100 ,1))
for i in range(row_sum):
    AY.ravel()[i::AY.shape[1]+row_sum] = 1

AY = AY.T

index0 = np.where(AY == 0)
index1 = np.where(AY == 1)

print(AX[index0,:])
print(AX[index1,:])