选择数组中向量的分量

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]]