Python numpy put_along_axis用法及代碼示例

本文簡要介紹 python 語言中 numpy.put_along_axis 的用法。

用法:

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

通過匹配一維索引和數據切片將值放入目標數組。

這會遍曆索引和數據數組中沿指定軸定向的匹配一維切片，並使用前者將值放入後者。這些切片可以是不同的長度。

沿軸返回索引的函數(如 argsort 和 argpartition )會為此函數生成合適的索引。

參數：

arr： ndarray (Ni…, M, Nk…)

目標數組。

indices： ndarray (Ni…, J, Nk…)

沿 arr 的每個 1d 切片更改的索引。這必須與 arr 的維度匹配，但 Ni 和 Nj 中的維度可能為 1 以針對 arr 進行廣播。

values： 數組 (Ni…, J, Nk…)

在這些索引處插入的值。它的形狀和尺寸被廣播以匹配 indices 。

axis： int

沿 1d 切片的軸。如果axis為None，則將目標數組視為已為其創建了扁平的一維視圖。

注意：

這相當於(但比)以下使用 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

例子：

對於這個示例數組

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

我們可以將最大值替換為：

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