如何在Pandas中基于多列多条件查找相似数据点？

数据集说明 #

你拥有一个包含DATE_TIME、ID、VALUE1、VALUE2、VALUE3、VALUE4、MODEL、SOLD等列的数据集，生成代码如下：

import numpy as np
import pandas as pd
import random

df = pd.DataFrame({'DATE_TIME': pd.date_range('2022-11-01', '2022-11-06 23:00:00', freq='20min'),
                   'ID': [random.randrange(1, 20) for n in range(430)]})

df['VALUE1'] = [random.randrange(110, 140) for n in range(430)]
df['VALUE2'] = [random.randrange(50, 60) for n in range(430)]
df['VALUE3'] = [random.randrange(80, 100) for n in range(430)]
df['VALUE4'] = [random.randrange(30, 50) for n in range(430)]

df['MODEL'] = [random.randrange(1, 3) for n in range(430)]

df['SOLD'] = [random.randrange(0, 2) for n in range(430)]

df['INSPECTION'] = df['DATE_TIME'].dt.day

df['MODE'] = np.select([df['INSPECTION'] == 1, df['INSPECTION'].isin([2, 3])], ['A', 'B'], 'C')

df['TIME'] = df['DATE_TIME'].dt.time
df['TIME'] = df['TIME'].astype('str')

def cycle_day_period(dataframe: pd.DataFrame, midnight='00:00:00', start_of_morning='06:00:00',
                     start_of_afternoon='13:00:00',
                     start_of_evening='18:00:00', end_of_evening='23:00:00', start_of_night='24:00:00'):
    bins = [midnight, start_of_morning, start_of_afternoon, start_of_evening, end_of_evening, start_of_night]
    labels = ['Night', 'Morning', 'Morning', 'Night', 'Night']

    return pd.cut(
        pd.to_timedelta(dataframe),
        bins=list(map(pd.Timedelta, bins)),
        labels=labels, right=False, ordered=False
    )

df['CYCLE_PART'] = cycle_day_period(df['TIME'], '00:00:00', '06:00:00', '13:00:00', '18:00:00', '23:00:00', '24:00:00')

需求目标 #

找到满足以下条件的数据点：

• VALUE1、VALUE2、VALUE3、VALUE4的值最相似（或相同）
• MODEL列值相同
• SOLD列值不同（一个为0，一个为1）

解决方案 #

我们可以通过按MODEL分组+计算特征相似度的方式实现，这里用欧氏距离衡量特征的相似性（距离越小，特征越相似）。

方法1：找到每个MODEL下最相似的记录对 #

# 定义需要计算相似度的特征列
value_cols = ['VALUE1', 'VALUE2', 'VALUE3', 'VALUE4']
result_list = []

# 按MODEL分组处理
for model, group in df.groupby('MODEL'):
    # 拆分SOLD=0和SOLD=1的子集
    sold0 = group[group['SOLD'] == 0].copy().reset_index()
    sold1 = group[group['SOLD'] == 1].copy().reset_index()
    
    # 如果该MODEL下没有两种SOLD类型，跳过
    if len(sold0) == 0 or len(sold1) == 0:
        continue
    
    # 向量化计算所有记录对的欧氏距离，提升效率
    sold0_vals = sold0[value_cols].values
    sold1_vals = sold1[value_cols].values
    distances = np.sqrt(((sold0_vals[:, np.newaxis] - sold1_vals)**2).sum(axis=2))
    
    # 找到距离最小的记录对索引
    min_dist_idx = np.unravel_index(np.argmin(distances), distances.shape)
    idx0 = sold0.loc[min_dist_idx[0], 'index']
    idx1 = sold1.loc[min_dist_idx[1], 'index']
    
    # 将符合条件的记录加入结果列表
    result_list.append(df.loc[idx0])
    result_list.append(df.loc[idx1])

# 转换为结果DataFrame并展示关键列
result_df = pd.DataFrame(result_list)
print(result_df[['ID', 'VALUE1', 'VALUE2', 'VALUE3', 'VALUE4', 'MODEL', 'SOLD']])

方法2：筛选所有满足相似度阈值的记录对 #

如果需要找出所有相似度足够高的记录对（而非仅最相似的一对），可以设定距离阈值筛选：

# 定义相似度阈值（可根据数据分布调整）
similarity_threshold = 10
value_cols = ['VALUE1', 'VALUE2', 'VALUE3', 'VALUE4']
result_list = []

for model, group in df.groupby('MODEL'):
    sold0 = group[group['SOLD'] == 0].copy().reset_index()
    sold1 = group[group['SOLD'] == 1].copy().reset_index()
    
    if len(sold0) == 0 or len(sold1) == 0:
        continue
    
    sold0_vals = sold0[value_cols].values
    sold1_vals = sold1[value_cols].values
    distances = np.sqrt(((sold0_vals[:, np.newaxis] - sold1_vals)**2).sum(axis=2))
    
    # 找到所有距离小于阈值的记录对
    pairs = np.where(distances < similarity_threshold)
    for i, j in zip(pairs[0], pairs[1]):
        idx0 = sold0.loc[i, 'index']
        idx1 = sold1.loc[j, 'index']
        result_list.append(df.loc[idx0])
        result_list.append(df.loc[idx1])

# 去重并转换为结果DataFrame
result_df = pd.DataFrame(result_list).drop_duplicates()
print(result_df[['ID', 'VALUE1', 'VALUE2', 'VALUE3', 'VALUE4', 'MODEL', 'SOLD']])

说明 #

• 欧氏距离是衡量数值型特征相似性的常用方法，你也可以根据需求替换为其他相似度指标（如曼哈顿距离）
• 阈值similarity_threshold需要根据你的数据分布调整，比如如果VALUE的取值范围是几十，阈值设10-15较为合理
• 代码中区分了MODEL和MODE列，确保按照需求中的MODEL列分组

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