numpy.put_along_axis

numpy.put_along_axis(arr, indices, values, axis)

通过匹配 1d 索引和数据切片将值放入目标数组中.

这会迭代索引和数据数组中沿指定轴定向的匹配的 1d 切片,并使用前者将值放入后者.这些切片可以具有不同的长度.

返回沿轴的索引的函数,如 argsort 和 argpartition ,会生成适合此函数的索引.

参数:

arrndarray (Ni…, M, Nk…)

目标数组.

indicesndarray (Ni…, J, Nk…)

要更改沿 arr 的每个 1d 切片的索引.这必须与 arr 的维度匹配,但 Ni 和 Nj 中的维度可以是 1,以针对 arr 进行广播.

valuesarray_like (Ni…, J, Nk…)

要在这些索引处插入的值.它的形状和维度被广播以匹配 indices 的形状和维度.

轴int

沿其取 1d 切片的轴.如果 axis 为 None,则目标数组被视为已创建了它的展平的 1d 视图.

参见

take_along_axis

通过匹配 1d 索引和数据切片从输入数组中获取值

注释

这等效于(但比以下使用 ndindex 和 s_ 更快),它将每个 ii 和 kk 设置为索引元组:

Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
J = indices.shape[axis]  # Need not equal M

for ii in ndindex(Ni):
    for kk in ndindex(Nk):
        a_1d       = a      [ii + s_[:,] + kk]
        indices_1d = indices[ii + s_[:,] + kk]
        values_1d  = values [ii + s_[:,] + kk]
        for j in range(J):
            a_1d[indices_1d[j]] = values_1d[j]

等效地,消除内部循环,最后两行将是:

a_1d[indices_1d] = values_1d

示例

Try it in your browser!

>>> import numpy as np

对于此示例数组

>>> a = np.array([[10, 30, 20], [60, 40, 50]])

我们可以用以下内容替换最大值:

>>> ai = np.argmax(a, axis=1, keepdims=True)
>>> ai
array([[1],
       [0]])
>>> np.put_along_axis(a, ai, 99, axis=1)
>>> a
array([[10, 99, 20],
       [99, 40, 50]])

Go BackOpen In Tab