任務是根據兩個相應向量的乘法結果逐行組合兩個陣列（構造排列）。如：
第 1_A 行、第 2_A 行、第 3_A 行、
第 1_B 行、第 2_B 行、第 3_B 行、
結果應該是：Row1_A_Row1_B、Row1_A_Row2_B、Row1_A_Row3_B、Row2_A_Row1_B等。
給定以下初始陣列：
n_rows = 1000
A = np.random.randint(10, size=(n_rows, 5))
B = np.random.randint(10, size=(n_rows, 5))
P_A = np.random.rand(n_rows, 1)
P_B = np.random.rand(n_rows, 1)
陣列 P_A 和 P_B 是各個陣列的對應向量，其中包含一個浮點數。如果乘法超過某個閾值，則組合行應僅出現在最終陣列中，例如：
lim = 0.8
我已經想到了以下功能或方法來解決這個問題，但我會對更快的解決方案感興趣。我愿意使用 numba 或其他庫，但理想情況下我想使用 numpy 改進矢量化解決方案。
方法一
def concatenate_per_row(A, B):
m1,n1 = A.shape
m2,n2 = B.shape
out = np.zeros((m1,m2,n1 n2),dtype=A.dtype)
out[:,:,:n1] = A[:,None,:]
out[:,:,n1:] = B
return out.reshape(m1*m2,-1)
%%timeit
A_B = concatenate_per_row(A, B)
P_A_B = (P_A[:, None]*P_B[None, :])
P_A_B = P_A_B.flatten()
idx = P_A_B > lim
A_B = A_B[idx, :]
P_A_B = P_A_B[idx]
每個回圈 37.8 ms ± 660 μs（7 次運行的平均值 ± 標準偏差，每次 10 次回圈）
方法B
%%timeit
A_B = []
P_A_B = []
for i in range(len(P_A)):
P_A_B_i = P_A[i]*P_B
idx = np.where(P_A_B_i > lim)[0]
if len(idx) > 0:
P_A_B.append(P_A_B_i[idx])
A_B_i = np.zeros((len(idx), A.shape[1] B.shape[1]), dtype='int')
A_B_i[:, :A.shape[1]] = A[i]
A_B_i[:, A.shape[1]:] = B[idx, :]
A_B.append(A_B_i)
A_B = np.concatenate(A_B)
P_A_B = np.concatenate(P_A_B)
每個回圈 9.65 ms ± 291 μs（7 次運行的平均值 ± 標準偏差，每次 100 次回圈）
uj5u.com熱心網友回復：
首先，有一個更高效的演算法。實際上，您可以預先計算輸出陣列的大小，以便可以將值直接寫入最終輸出陣列中，而不是臨時存盤在串列中。為了有效地找到大小，您可以對陣列進行排序P_B，然后進行二分搜索，以找到大于lim/P_A[i,0]所有可能的值的數量i（P_B*P_A[i,0] > lim相當于P_B > lim/P_A[i,0]）。i可以臨時存盤為每個過濾的專案數，以便快速回圈過濾專案。
此外，您可以使用Numba顯著加快主回圈的計算速度。
這是結果代碼：
@nb.njit('(int_[:,::1], int_[:,::1], float64[:,::1], float64[:,::1])')
def compute(A, B, P_A, P_B):
assert P_A.shape[1] == 1
assert P_B.shape[1] == 1
P_B_sorted = np.sort(P_B.reshape(P_B.size))
counts = len(P_B) - np.searchsorted(P_B_sorted, lim/P_A[:,0], side='right')
n = np.sum(counts)
mA, mB = A.shape[1], B.shape[1]
m = mA mB
A_B = np.empty((n, m), dtype=np.int_)
P_A_B = np.empty((n, 1), dtype=np.float64)
k = 0
for i in range(P_A.shape[0]):
if counts[i] > 0:
idx = np.where(P_B > lim/P_A[i, 0])[0]
assert counts[i] == len(idx)
start, end = k, k counts[i]
A_B[start:end, :mA] = A[i, :]
A_B[start:end, mA:] = B[idx, :]
P_A_B[start:end, :] = P_B[idx, :] * P_A[i, 0]
k = counts[i]
return A_B, P_A_B
以下是我機器上的性能結果：
Original: 35.6 ms
Optimized original: 18.2 ms
Proposed (with order): 0.9 ms
Proposed (no ordering): 0.3 ms
上面提出的演算法比原始優化演算法快 20 倍。它可以做得更快。事實上，如果專案的順序無關緊要，您可以使用argsortso 對B和重新排序P_B。這使你不來計算idx每次在熱回圈，并選擇直接從最后一個元素B和P_B（即保證是在相同的順序比原來的代碼高于閾值，但沒有）。由于所選專案連續存盤在記憶體中，因此此實作要快得多。最后，最后一個實作快了大約60 倍比原來的優化演算法。請注意，即使沒有 Numba，建議的實作也比原始實作快得多。
這是不關心專案順序的實作：
@nb.njit('(int_[:,::1], int_[:,::1], float64[:,::1], float64[:,::1])')
def compute(A, B, P_A, P_B):
assert P_A.shape[1] == 1
assert P_B.shape[1] == 1
nA, mA = A.shape
nB, mB = B.shape
m = mA mB
order = np.argsort(P_B.reshape(nB))
P_B_sorted = P_B[order, :]
B_sorted = B[order, :]
counts = nB - np.searchsorted(P_B_sorted.reshape(nB), lim/P_A[:,0], side='right')
nRes = np.sum(counts)
A_B = np.empty((nRes, m), dtype=np.int_)
P_A_B = np.empty((nRes, 1), dtype=np.float64)
k = 0
for i in range(P_A.shape[0]):
if counts[i] > 0:
start, end = k, k counts[i]
A_B[start:end, :mA] = A[i, :]
A_B[start:end, mA:] = B_sorted[nB-counts[i]:, :]
P_A_B[start:end, :] = P_B_sorted[nB-counts[i]:, :] * P_A[i, 0]
k = counts[i]
return A_B, P_A_B
轉載請註明出處，本文鏈接：https://www.uj5u.com/shujuku/398140.html