如何为 3D numpy 数组中的不同行分配不同数量的元素？

Question

我正在尝试使用一个数组来索引另一个数组。目标是将值 (val) 在第一行（第二轴）中放置不同的时间数（num）。我不太擅长解释这一点，但示例的结果显示了我正在寻找的内容。我想我已经很接近了。

我收到此错误：

TypeError：只有整数标量数组可以转换为标量索引

import numpy as np 

# dataset
data = np.zeros((3, 4, 5))
val = np.array([6, 7, 8])
num = np.array([2, 4, 3])

# Do something like this, using an index, not a loop
data[:, 0][:, 0:num] = val

# This is the result I am hoping to get

[[[6. 6. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[7. 7. 7. 7. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[8. 8. 8. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]]

Answer 1

假设data 实际上全为零，您可以使用np.cumsum 使用我通常用于屏蔽的技术非常简单地做到这一点。关键是使用简单的索引设置运行的开始和结束，然后在它们之间求和。

data[:, 0, 0] = val   # Start the sum
data[np.arange(data.shape[0]), 0, num] = -val
data.cumsum(axis=-1, out=data)

如果num == data.shape[-1] 曾经是True，则需要先过滤索引，然后求和直到结束：

mask = (num < data.shape[-1])
data[np.arange(data.shape[0])[mask], 0, num[mask]] = -val

如果data 不是全零，您可以轻松创建一个掩码并将其与np.repeat 结合使用。您使用与np.cumsum 相同的技术以及np.uint8 和np.bool_ 具有相同项目大小的事实来制作掩码：

mask = np.zeros_like(data, dtype=np.uint8)
mask[:, 0, 0] = 1   # Start the sum
m = (num < data.shape[-1])
mask[np.arange(data.shape[0])[m], 0, num[m]] = -1
mask.cumsum(axis=-1, out=mask)
mask = mask.view(bool)

data[mask] = np.repeat(val, num)

如果您只对每个平面的第一行感兴趣，您可能会在仅包含这些行的视图上进行操作以简化索引：

x = data[:, 0, :]

Answer 2

使用以下代码：

ind = (np.repeat(range(3), num),
       np.zeros(num.sum(), dtype='int'),
       np.hstack([np.arange(i) for i in num]))
tVal = np.repeat(val, num)
data[ind] = tVal

打印 ind 和 tVal 以查看中间对象。

详情：

1. ind 是索引数组的元组（对于每个维度）。首先 第一个维度的数组，依此类推。它是一个“复合索引” 要设置的元素。

2. tVal 包含要在 data 的连续元素中设置的值。

3. data[ind] = tVal 执行实际替换（来自 tVal 的每个值 到 data 的指定元素中）。

阅读每个使用的 Numpy 方法以了解其他细节 和背景。