基于多列连接两个NUMPY数组：实现内连接与左连接

原始数组 #

import numpy as np

arr1 = np.array([['a', 1, 23],
                 ['a', 1, 35],
                 ['b', 2, 27],
                 ['b', 3, 98]], dtype=object)

arr2 = np.array([['a', 1, 'g'],
                 ['b', 3, 'h']], dtype=object)

需求说明 #

按数组的第0、1列作为连接键，分别实现内连接和左连接，无需使用Pandas。

1. 内连接实现 #

内连接仅保留两个数组中键匹配的行，将arr1的行与arr2对应键的第三列值拼接。

# 提取两个数组的连接键（第0、1列）
keys1 = arr1[:, :2]
keys2 = arr2[:, :2]

# 将键转换为元组，方便后续匹配比较
keys1_tuple = [tuple(k) for k in keys1]
keys2_tuple = [tuple(k) for k in keys2]

# 筛选arr1中键存在于arr2的行索引
inner_match_indices = [i for i, key in enumerate(keys1_tuple) if key in keys2_tuple]
# 获取对应匹配的arr2第三列值
inner_match_values = [arr2[keys2_tuple.index(key), 2] for key in keys1_tuple if key in keys2_tuple]

# 拼接得到内连接结果
arr_inner = np.column_stack((arr1[inner_match_indices], inner_match_values))

输出结果：

[['a' 1 23 'g']
 ['a' 1 35 'g']
 ['b' 3 98 'h']]

2. 左连接实现 #

左连接保留arr1的所有行，匹配到arr2键的行拼接对应值，未匹配的行填充NaN。

# 遍历arr1的每个键，匹配arr2的值或填充NaN
left_match_values = []
for key in keys1_tuple:
    if key in keys2_tuple:
        left_match_values.append(arr2[keys2_tuple.index(key), 2])
    else:
        left_match_values.append(np.nan)

# 拼接得到左连接结果
arr_left = np.column_stack((arr1, left_match_values))

输出结果：

[['a' 1 23 'g']
 ['a' 1 35 'g']
 ['b' 2 27 nan]
 ['b' 3 98 'h']]

内容的提问来源于stack exchange，提问作者user11809571