Question

我想编写一个函数amap，它类似于内置map，但支持numpy广播和输出数组。它会像这样工作：

>>> amap(lambda x:x**2, 1)
array(1)
>>> amap(lambda x:x**2, [1, 2])
array([1, 4])
>>> amap(lambda x,y:y**2+x**2, 1, [1, 2])
array([2, 5])
>>> amap(lambda x,y:y**2+x**2, [1, 2], [[1], [2]])
array([[2, 5], [5, 8]])

传入的函数在标量上运行。

我写了一个实现。

def amap(func, *args):
    '''array version of build-in map
    amap(function, scalar/sequence[, ...]) -> array
    '''
    args = np.broadcast(None, *args)
    res = [func(*arg[1:]) for arg in args]
    res = np.asarray(res).reshape(args.shape)
    return res

它有效，但它有些丑陋而且效率不高。你有什么好主意吗？特别是那些更时尚的东西？

Answer 1

直接使用numpy.vectorize？

>>> numpy.vectorize(lambda x:x**2)(1)
array(1)
>>> numpy.vectorize(lambda x:x**2)([1,2])
array([1, 4])
>>> numpy.vectorize(lambda x, y: y**2 + x**2)(1, [1,2])
array([2, 5])
>>> numpy.vectorize(lambda x, y: y**2 + x**2)([1,2], [[1], [2]])
array([[2, 5],
       [5, 8]])

功能类似于内置映射但支持numpy广播

1 个答案: