脾气暴躁的怪癖:将函数应用于两个1D数组的所有对，以获得一个2D数组

问题描述

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

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

提前谢谢！

解决方案

您可以使用 numpy广播对两个数组进行计算，使用newaxis将a转换为垂直2D数组:

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]])

要进行基准测试:

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

Let's say I have 2 one-dimensional (1D) numpy arrays, a and b, with lengths n1 and n2 respectively. I also have a function, F(x,y), that takes two values. Now I want to apply that function to each pair of values from my two 1D arrays, so the result would be a 2D numpy array with shape n1, n2. The i, j element of the two-dimensional array would be F(a[i], b[j]).

I haven't been able to find a way of doing this without a horrible amount of for-loops, and I'm sure there's a much simpler (and faster!) way of doing this in numpy.

Thanks in advance!

解决方案

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数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！