numpy.mask_indices#

numpy.mask_indices(n, mask_func, k=0)[source]#

给定一个掩码函数,返回用于访问 (n, n) 数组的索引。

假设 mask_func 是一个函数,对于大小为 (n, n) 的方阵 a,当以 mask_func(a, k) 的形式调用时(可能带有偏移参数 k),它会返回一个在新位置包含零的新数组(如 triutril 等函数就是如此)。然后,此函数返回非零值所在的索引。

参数:
nint

返回的索引将有效用于访问形状为 (n, n) 的数组。

mask_funccallable

一个其调用签名类似于 triu, tril 的函数。也就是说,mask_func(x, k) 返回一个形状与 x 相同的布尔数组。k 是该函数的一个可选参数。

kscalar

一个可选参数,它会传递给 mask_func。像 triu, tril 这样的函数会接受一个被解释为偏移量的第二个参数。

返回:
indices数组元组。

mask_func(np.ones((n, n)), k) 为 True 的位置相对应的 n 个索引数组。

示例

>>> import numpy as np

这些索引将允许您访问任何 3x3 数组的上三角部分

>>> iu = np.mask_indices(3, np.triu)

例如,如果 a 是一个 3x3 数组

>>> a = np.arange(9).reshape(3, 3)
>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> a[iu]
array([0, 1, 2, 4, 5, 8])

偏移量也可以传递给掩码函数。这会给我们主对角线右侧第一个对角线开始的索引

>>> iu1 = np.mask_indices(3, np.triu, 1)

现在我们只提取三个元素

>>> a[iu1]
array([1, 2, 5])