python - 有没有办法使用矢量化操作具有 1D 的 3D 数组？

Question

我正在尝试在 1 轴的前 n 个位置创建一个 3D 数组，其中 n 来自长度为 0 轴的数组。我正在处理一个大型数据集并试图加快速度。
我认为代码会更有意义。我正在尝试对 for 循环进行矢量化。

    import numpy as np 

    data = np.zeros((3, 4, 5))
    num = np.array([2, 4, 3])

    for i in range(len(num)):
        data[:, 0][i, 0:num[i]] = 1

    print(data)

这是我正在寻找的结果:

[[[1. 1. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[1. 1. 1. 1. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[1. 1. 1. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]]

Answer 1

这有点hacky但有效:

>>> data = np.zeros((3, 4, 5))
>>> num = np.array([2, 4, 3])
>>> max_idx = min(data.shape[2], np.max(num))
>>> data[:,0,:max_idx] = np.arange(max_idx) < num[:,np.newaxis]
>>> data
array([[[1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

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

       [[1., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]]])

它基于以下想法:

>>> np.arange(5)
array([0, 1, 2, 3, 4])
>>> (np.arange(5) < 2).astype(np.float64)
array([1., 1., 0., 0., 0.])

除了 2替换为垂直广播 num生成“前 N 个元素设置为 1，其余为 0”。