python - 如何找到特定数字的最大索引和最小索引列明智的numpy

Question

我有一个像

这样的 numpy 数组

m = np.array([[0,0,0,0,0],
              [0,0,0,1,1],
              [0,1,0,1,0],
              [0,1,0,1,0],
              [0,0,1,1,1],])

(m == 1).argmax(0) 将给出 array([0, 2, 4, 1, 1])。是否有任何类似的函数来获取每列中 1 的最小索引和最大索引。即

array([[ nan, 2., 4., 1., 1.], [ nan, 3., 4., 4., 4.]])

Answer 1

一种方法是——

mask = m==1
mask_cumsum = mask.cumsum(0)

valid_mask = mask.any(0)
min_idx = (mask_cumsum==1).argmax(0)
max_idx = mask_cumsum.argmax(0)
min_max_idx = np.vstack((min_idx, max_idx))

out = np.where(valid_mask, min_max_idx, np.nan)