Python:3d numpy数组中的高效搜索

2024-03-29 04:54:31 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个3d numpy数组(shape (z,y,x): 137,601,1200)),面临着在每个垂直列中找到最接近500的值的索引的挑战。 例如,我想为数组(:,30,112)找到最接近500的索引。数字从下到上按降序排列。例如,(:,30,112)的137个值中的一些可以如下所示:[1033.91 1031.35 ... 0.01]。你知道吗

我的第一种方法是通过for循环对每个垂直列应用搜索函数。但我相信,一定有更快的办法!这是我的第一次尝试:

from bisect import bisect_left

def takeClosest(myList, myNumber):

    pos = bisect_left(myList, myNumber)
    if pos == 0:
        return 0
    if pos == len(myList):
        return -1
    before = myList[pos - 1]
    after = myList[pos]
    if after - myNumber < myNumber - before:
       return pos
    else:
       return pos-1


def calcOptK(tnsr):
    # 2d array which saves the optimal levels
    optKs = np.zeros([601,1200])
    #iterate through the tensor and call on each point the search function    
    for y in range (0, 601):
        for x in range (0, 1200):
            optKs[y,x]=136-takeClosest(p3d[:,y,x][::-1],500)
    return optKs; 

optKs=calcOptK(p3d)

更新:感谢@hpaulj和@Mstaino。 这非常有效:

optKs = np.argmin(np.abs(p3d-500), axis=0)

Tags: theposforreturnifdefnp数组