查找 NumPy 数组中大于 N 的值的开始/停止索引范围

sla*_*law 5 python numpy

假设我有一个 NumPy 数组:

x = np.array([2, 3, 4, 0, 0, 1, 1, 4, 6, 5, 8, 9, 9, 4, 2, 0, 3])
Run Code Online (Sandbox Code Playgroud)

对于 中的所有值x >= 2,我需要找到连续值x >=2(即,一个大于或等于 2 的单个值的运行不计算在内)的开始/停止索引。然后,我对x >= 3, x >=4, ...,重复此操作x >= x.max()输出应该是一个三列的 NumPy 数组(第一列是最小值,第二列是包含开始索引,第三列是停止索引),看起来像:

[[2,  0,  2],
 [2,  7, 14],
 [3,  1,  2],
 [3,  7, 13],
 [4,  7, 13],
 [5,  8, 12],
 [6, 10, 12],
 [8, 10, 12],
 [9, 11, 12]
]
Run Code Online (Sandbox Code Playgroud)

天真地,我可以查看每个唯一值,然后搜索开始/停止索引。但是,这需要对 进行多次传递x。完成此任务的最佳 NumPy 矢量化方法是什么?是否有不需要多次传递数据的解决方案?

更新

我意识到我还需要计算单个实例。所以,我的输出应该是:

[[2,  0,  2],
 [2,  7, 14],
 [2, 16, 16],  # New line needed
 [3,  1,  2],
 [3,  7, 13],
 [3, 16, 16],  # New line needed
 [4,  2,  2],  # New line needed
 [4,  7, 13],
 [5,  8, 12],
 [6,  8,  8],  # New line needed
 [6, 10, 12],
 [8, 10, 12],
 [9, 11, 12]
]
Run Code Online (Sandbox Code Playgroud)

And*_*dyK 3

这是另一个解决方案(我相信可以改进):

import numpy as np
from numpy.lib.stride_tricks import as_strided

x = np.array([2, 3, 4, 0, 0, 1, 1, 4, 6, 5, 8, 9, 9, 4, 2, 0, 3])

# array of unique values of x bigger than 1
a = np.unique(x[x>=2])

step = len(a)  # if you encounter memory problems, try a smaller step
result = []
for i in range(0, len(a), step):
    ai = a[i:i + step]
    c = np.argwhere(x >= ai[:, None])
    c[:,0] = ai[c[:,0]]
    c =  np.pad(c, ((1,1), (0,0)), 'symmetric')

    d = np.where(np.diff(c[:,1]) !=1)[0]

    e = as_strided(d, shape=(len(d)-1, 2), strides=d.strides*2).copy()
    # e = e[(np.diff(e, axis=1) > 1).flatten()]
    e[:,0] = e[:,0] + 1 

    result.append(np.hstack([c[:,0][e[:,0, None]], c[:,1][e]]))

result = np.concatenate(result)

# array([[ 2,  0,  2],
#        [ 2,  7, 14],
#        [ 2, 16, 16],
#        [ 3,  1,  2],
#        [ 3,  7, 13],
#        [ 3, 16, 16],
#        [ 4,  2,  2],
#        [ 4,  7, 13],
#        [ 5,  8, 12],
#        [ 6,  8,  8],
#        [ 6, 10, 12],
#        [ 8, 10, 12],
#        [ 9, 11, 12]])
Run Code Online (Sandbox Code Playgroud)

抱歉没有评论每个步骤的作用——如果以后我有时间我会修复它。