python - numpy 选择不同长度的每一行

Question

我有两个数组

x=array([[0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [2, 2, 2, 2, 2]])

我想按数组 y 中的长度对每一行中的元素进行子选择

y = array([3, 2, 4])

我的目标是 z:

z = array([[0, 0, 0],
   [1, 0,],
   [2, 2, 2, 2]])

我如何使用 numpy 函数而不是列表/循环来做到这一点？

非常感谢您的帮助。

Answer 1

Numpy 数组针对具有特定维度的同构数组进行了优化。我喜欢把它想象成一个矩阵:在每行上有不同数量的元素是没有意义的。

也就是说，根据你想如何使用处理后的数组，你可以简单地制作一个数组列表:

z = [array([0, 0, 0]),
array([1, 0,]),
array([2, 2, 2, 2]])]

不过，您仍需要手动执行此操作:

x = array([[0, 0, 0, 0, 0], [1, 0, 0, 0, 0], [2, 2, 2, 2, 2]])
y = array([3, 2, 4])

z = [x_item[:y_item] for x_item, y_item in zip(x, y)]

列表理解遍历 x 和 y 并结合 zip() 创建原始数组的新切片。