Python：结合两个map（）逻辑

Question

在我的代码中，map()函数的f(a,b)需要两个1x3的numpy数组作为输入。因为这是从C代码转换过来的别人的实现，所以我不能改变f()。

第二个输入y总是一个Nx3的numpy数组。第一个输入x有两种情况。

在一种情况下，它是一个1x3的numpy数组，因此我执行以下操作：

unwrap = partial(f, x)
result = map(unwrap, y)

在另一种情况下，它是一个Nx3的numpy数组，然后我执行

unwrap = f
result = map(unwrap, x, y)

有没有一种方法可以将这两种情况合并在一起？

- nos

Answer 1

在第二种情况下，您正在将x和y进行压缩，并将其传递给f。在第一种情况下，您正在将f应用于x并将其映射到y上 - 换句话说，您正在将f应用于一个固定参数和一个变量参数。

因此，您应该能够看到将f应用于常数参数与将其映射到常数列表上是相同的。特别地，

map(partial(f, x), y) == map(f, [x]*len(y), y)

Aside: Haskellers will recognise this as one of the applicative laws:

pure (f x) = fmap f (pure x)

虽然它们的行为相同，但它们并没有完全相同的特点。我预计第二个会分配更多的内存并且运行时间更长。但这只是微小的优化领域；除非你已经对其进行了性能分析，并确定你确实需要这段代码尽可能地快，否则最好选择更清晰的选项。

- Benjamin Hodgson

Answer 2

np.broadcast_to 可以将 A '重塑' 以匹配 B; 然后你可以同时迭代这两个数组。它使用 striding，因此实际上不会增加内存使用。

In [370]: def f(a,b):
     ...:     assert(a.shape==(1,3))
     ...:     assert(b.shape==(1,3))
     ...:     return a+b
     ...: 
In [371]: B=np.arange(12).reshape(4,3)
In [372]: A=np.arange(3).reshape(1,3)
In [373]: np.broadcast_to(A, B.shape)   # (1,3) to (4,3)
Out[373]: 
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])
In [374]: np.broadcast_to(B, B.shape)   # no change with (4,3)
Out[374]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

我通常使用列表推导式而不是 map 函数：

In [375]: [f(np.atleast_2d(a),np.atleast_2d(b)) for a,b in zip(np.broadcast_to(A,B.shape),B)]
Out[375]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

In [376]: [f(np.atleast_2d(a),np.atleast_2d(b)) for a,b in zip(np.broadcast_to(B,B.shape),B)]
Out[376]: 
[array([[0, 2, 4]]),
 array([[ 6,  8, 10]]),
 array([[12, 14, 16]]),
 array([[18, 20, 22]])]

迭代2D数组会产生1D数组的列表，因此需要使用np.atleast_2d来满足我的f断言。如果f也接受（3，）输入，则不需要这样做。或者使用map：

In [377]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), np.broadcast_to(B,B.shape),B)
Out[377]: <map at 0xb14f4c6c>
In [378]: list(_)
Out[378]: 
[array([[0, 2, 4]]),
 array([[ 6,  8, 10]]),
 array([[12, 14, 16]]),
 array([[18, 20, 22]])]
In [379]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), np.broadcast_to(A,B.shape),B)
Out[379]: <map at 0xb0871a8c>
In [380]: list(_)
Out[380]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

np.vectorize和np.frompyfunc也可以处理这种广播，但它们是为接受标量的函数设计的，而不是1d数组。

使用broadcast_arrays，我可以平等地处理两个数组：

In [386]: map(lambda a,b: f(np.atleast_2d(a),np.atleast_2d(b)), *np.broadcast_arrays(B,A))
Out[386]: <map at 0xb69851ac>
In [387]: list(_)
Out[387]: 
[array([[0, 2, 4]]),
 array([[3, 5, 7]]),
 array([[ 6,  8, 10]]),
 array([[ 9, 11, 13]])]

更一般地说，A和B可以是任何产生所需(N,3)数组的东西。我可以通过生成(N,1,3)数组来摆脱使用atleast_2d：

In [397]: map(f, *np.broadcast_arrays(np.arange(3)[None,None,:], np.arange(0,40,10)[:,None,None]))
Out[397]: <map at 0xb08b562c>
In [398]: list(_)
Out[398]: 
[array([[0, 1, 2]]),
 array([[10, 11, 12]]),
 array([[20, 21, 22]]),
 array([[30, 31, 32]])]

- hpaulj