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