python - 选择数组中向量的分量

Question

我有一个形状为 (7, 2, 3) 的张量
我想从 7 个 2x3 矩阵中的每一个中选择两个行向量之一，即

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

 [[0, 0, 0],
  [1, 1, 1]],
 ...x7
]

到

a = [
 [0, 0, 0],
 [0, 0, 0],
 [0, 0, 0]
 ...x7
]
b = [
 [1, 1, 1],
 [1, 1, 1],
 [1, 1, 1]
 ...x7
]

每个都有形状 (7, 3)。
如果没有 reshape 我怎么能做到这一点? (当某些尺寸相同时，我发现 reshape 有点令人困惑)。
我也知道

np.array(map(lambda item: item[0], x)))

但如果有的话，我想要一种更简洁的方式。

Answer 1

只需使用循环索引:data[:, i, :]哪里i从 0 到 1 的循环

import numpy as np

a = np.array([
 [[0, 0, 0],
  [1, 1, 1]],

 [[0, 0, 0],
  [1, 1, 1]]
])
print(a[:, 1, :])

会产生

[[1 1 1]
 [1 1 1]]