admin管理员组

文章数量:1410689

In Python, I have two arrays:

import numpy as np

a = np.array([
  [1, 0.2],
  [2, 0.5],
  [3, 0.8]])

b = np.array([
  [2, 0.4],
  [3, 0.7],
  [4, 1.3],
  [5, 2]])

I need to do an "outer join" based on the values of the first column, and fill with 0 in case there is no match, i.e., end up with something like this:

c = np.array([
  [1, 0.2, 0],
  [2, 0.5, 0.4],
  [3, 0.8, 1.7],
  [4, 0, 1.3],
  [5, 0, 2]])

I was first thinking about converting to DataFrames and then doing an outer join, but I need this to be very fast, since I need to do it 18 million times. I don't necessarily need the entire array, two vectors in the same order would also suffice. Having the two arrays as DataFrames in the first place is not an option, because I use the arrays for several operations where NumPy has considerable speed advantage.

How can I do this very fast?

In Python, I have two arrays:

import numpy as np

a = np.array([
  [1, 0.2],
  [2, 0.5],
  [3, 0.8]])

b = np.array([
  [2, 0.4],
  [3, 0.7],
  [4, 1.3],
  [5, 2]])

I need to do an "outer join" based on the values of the first column, and fill with 0 in case there is no match, i.e., end up with something like this:

c = np.array([
  [1, 0.2, 0],
  [2, 0.5, 0.4],
  [3, 0.8, 1.7],
  [4, 0, 1.3],
  [5, 0, 2]])

I was first thinking about converting to DataFrames and then doing an outer join, but I need this to be very fast, since I need to do it 18 million times. I don't necessarily need the entire array, two vectors in the same order would also suffice. Having the two arrays as DataFrames in the first place is not an option, because I use the arrays for several operations where NumPy has considerable speed advantage.

How can I do this very fast?

Share asked Mar 5 at 17:51 EmilAEmilA 1797 bronze badges
Add a comment  | 

2 Answers 2

Reset to default 1

How about something like this?

keys = np.union1d(a[:, 0], b[:, 0])

# Initialize result array with zeros - assume the 3 columns as output
c = np.zeros((keys.shape[0], 3))
c[:, 0] = keys

idx_a = np.searchsorted(keys, a[:, 0])
idx_b = np.searchsorted(keys, b[:, 0])

# Assign values where keys match
c[idx_a, 1] = a[:, 1]
c[idx_b, 2] = b[:, 1]

print(c)

"""
[[1.  0.2 0. ]
 [2.  0.5 0.4]
 [3.  0.8 0.7]
 [4.  0.  1.3]
 [5.  0.  2. ]]
"""

Another possible solution:

keys = np.union1d(a[:, 0], b[:, 0])

ind_a = np.equal.outer(keys, a[:, 0])
ind_b = np.equal.outer(keys, b[:, 0])

A = ind_a @ a[:, 1]
B = ind_b @ b[:, 1]

c = np.column_stack((keys, A, B))

This solution computes the union of keys from both arrays (np.union1d), generates indicator binary matrices (ind_a, ind_b) using outer equality checks to map original keys to the union (np.ufunc.outer), then multiplies them with the value columns to align values (unmatched entries become 0). Finally, it stacks keys and aligned values (np.column_stack).

Output:

[[1.  0.2 0. ]
 [2.  0.5 0.4]
 [3.  0.8 0.7]
 [4.  0.  1.3]
 [5.  0.  2. ]]

本文标签: pythonOuter join numpy arraysfill with 0Stack Overflow