python - numpy 沿第一个和最后一个出现值出现的轴查找切片

Question

我有一个带有整数值的 3D numpy 数组，定义如下:

import numpy as np
x = np.random.randint(0,  100, (10, 10, 10))

现在我想要做的是找到沿给定轴(比如 1)出现特定值的最后一个切片(或者第一个切片)。目前，我会做类似的事情:

first=None
last=None
val = 20

for i in range(len(x.shape[1]):
    slice = x[:, i, :]
    if len(slice[slice==val]) > 0:
        if not first:
            first = i
        last = i

return first, last

这似乎有点不符合 Python 风格，我想知道是否有一些 numpy 魔法可以做到这一点？

Answer 1

您可能可以对其进行优化以使其更快，但这是您搜索内容的矢量化版本:

axis = 1
mask = np.where(x==val)[axis]
first, last = np.amin(mask), np.amax(mask)

它首先使用 np.where 在您的数组中找到元素 val 并返回 min 和 max沿所需轴的索引。