为一维 numpy 数组创建成对二维数组的更有效方法是什么?

What is the more efficient way to create a pairwise 2D array for a 1D numpy array?

给定 2 个长度为 N 的 numPy 数组,我想基于自定义函数创建一个成对的二维数组 (N x N)。

    import numpy as np
    import pandas as pd
    
    A = np.array(A) # size N
    B = np.array(B) # size N
    Fij = f(A[i], B[i], A[j], B[j])
    
    # f is a pairwise function of Ai, Bi and Aj, Bj

我想创建一个大小为 NxN 的数组 C。这样

    C = Fij # with i, j in range(N)

示例:

    A = [1,2,3,4]
    B = ['a', 'b', 'c', 'd']
    
    def f(i,j):
        return str(A[i]+A[j])+B[i]+B[j]
    # f(0,1) = 3ab

我要创作

C = [[2aa, 3ab, 4ac, 5ad],
     [3ba, 4bb, 5bc, 6bd], 
     [4ca, 5cb, 6cc, 7cd], 
     [5da, 6db, 7dc, 8dd]]

我知道我可以通过嵌套循环来做到这一点:

C=np.empty([N, N])
for i in range(N):
    for j in range(N):
        C[i, j] = f(i,j)

但正在寻找更简洁、更高效的解决方案。

如果您无法分解函数的核心,您可以这样做:

import numpy as np
from itertools import product

A = [1,2,3,4]
B = ['a', 'b', 'c', 'd']
N = len(A)

def f(i,j):
    return str(A[i] + A[j]) + B[i] + B[j]

arr = np.array([f(i,j) for i, j in product(range(N), range(N))]).reshape(N,N)

在 numpy 中,您可以通过使用 + 来添加整数和连接字符来做出很好的解决方案。

A = np.array([1,2,3,4])
B = np.char.array(['a', 'b', 'c', 'd']) # char array so can use "+" instead of "add"

A0 = A[:,None] + A[None,:]  # add the integers
B0 = B[:,None] + B[None,:]  # concatenate the characters
C = A0.astype('<U1') + B0

或单行:

C = (A[:,None] + A[None,:]).astype('<U1') + B[:,None] + B[None,:]

打印出上面的每一行,得到:

A0
[[2 3 4 5]
 [3 4 5 6]
 [4 5 6 7]
 [5 6 7 8]]

B0
[['aa' 'ab' 'ac' 'ad']
 ['ba' 'bb' 'bc' 'bd']
 ['ca' 'cb' 'cc' 'cd']
 ['da' 'db' 'dc' 'dd']]

C
[['2aa' '3ab' '4ac' '5ad']
 ['3ba' '4bb' '5bc' '6bd']
 ['4ca' '5cb' '6cc' '7cd']
 ['5da' '6db' '7dc' '8dd']]