python - 向量化查找数组中最接近值的函数

Question

我编写了以下函数:

import numpy as np


def _find_nearest(array, value):
    """Find the index in array whose element is nearest to value.

    Parameters
    ----------
    array : np.array
      The array.

    value : number
      The value.

    Returns
    -------
    integer
      The index in array whose element is nearest to value.

    """
    if array.argmax() == array.size - 1 and value > array.max():
        return array.size
    return (np.abs(array - value)).argmin()

我想向量化这个函数，以便我可以一次传递多个值。也就是说，我想要 value是一个数组，并且有 _find_nearest返回提交的 value_array 中每个值的索引，而不是单个索引。 .

有人能找到一种方法吗？

Answer 1

在父函数内，值和数组均可见，您可以使用lambda来启用矢量化。我将调用父函数 main

def main():
    value = np.random.rand(10, 1)
    array = np.random.rand(100, 100)
    vec_nearest = lambda x: _find_nearest(array, x)
    np.vectorize(vec_nearest)(value)

这适用于一个数组和多个向量值。返回将是一个数组。