脾气暴躁的怪癖:将函数应用于两个1D数组的所有对，以获得一个2D数组 [英] Numpy quirk: Apply function to all pairs of two 1D arrays, to get one 2D array

问题描述

假设我有2个一维(1D)numpy数组a和b，其长度分别为n1和n2.我还有一个函数F(x,y)，它带有两个值.现在，我想将该函数应用于我的两个1D数组中的每对值，因此结果将是形状为n1, n2的2D numpy数组.二维数组的i, j元素为F(a[i], b[j]).

在没有大量for循环的情况下，我一直找不到找到这种方法的方法，而且我敢肯定，在numpy中，有一种更简单(而且更快！)的方法.

提前谢谢！

解决方案

您可以使用 解决方案

You can use numpy broadcasting to do calculation on the two arrays, turning a into a vertical 2D array using newaxis:

In [11]: a = np.array([1, 2, 3]) # n1 = 3
    ...: b = np.array([4, 5]) # n2 = 2
    ...: #if function is c(i, j) = a(i) + b(j)*2:
    ...: c = a[:, None] + b*2

In [12]: c
Out[12]: 
array([[ 9, 11],
       [10, 12],
       [11, 13]])

To benchmark:

In [28]: a = arange(100)

In [29]: b = arange(222)

In [30]: timeit r = np.array([[f(i, j) for j in b] for i in a])
10 loops, best of 3: 29.9 ms per loop

In [31]: timeit c = a[:, None] + b*2
10000 loops, best of 3: 71.6 us per loop

这篇关于脾气暴躁的怪癖:将函数应用于两个1D数组的所有对，以获得一个2D数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！