如何使查找关联行的Python函数兼容int/float类型输入？

原数据与问题函数 #

给定的DataFrame：

import pandas as pd
import numpy as np
from typing import Union

data = pd.DataFrame([
    {"id": 1, "user_id": 999, "phone_number": "61412308310", "email": "can@gmail.com"},
    {"id": 2, "user_id": 129, "phone_number": "61477708777", "email": "acdc@gmail.com"},
    {"id": 3, "user_id": 213, "phone_number": "61488908495", "email": "adel99@gmail.com"},
    {"id": 4, "user_id": 145, "phone_number": "61477708777", "email": "austr@gmail.com"},
    {"id": 5, "user_id": 214, "phone_number": "61421445777", "email": "austr@gmail.com"},
    {"id": 6, "user_id": 214, "phone_number": "61421445326", "email": "jango@gmail.com"},
])

原查找关联行的函数仅支持list类型输入，传入int/float时会报错——因为data.isin(value)要求输入是可迭代对象，单个数值不满足要求：

def get_related_values(data: pd.DataFrame, value: Union[int, list]) -> pd.DataFrame:
        
    related_values = data[data.isin(value).any(axis=1)]
    
    if set(np.array(value)) != set(np.array(related_values).reshape(-1)):
        return get_related_values(data, np.array(related_values).reshape(-1)) 
    else:
        return related_values

修改方案 #

只需在函数开头把单个的int/float值转换成包含该值的列表，让后续isin调用始终能接收可迭代对象，修改后的函数如下：

def get_related_values(data: pd.DataFrame, value: Union[int, float, list]) -> pd.DataFrame:
    # 统一输入格式：单个数值转成单元素列表
    if isinstance(value, (int, float)):
        value = [value]
    
    related_values = data[data.isin(value).any(axis=1)]
    
    # 展平结果并转成list，保证递归调用的输入格式统一
    flattened_vals = np.array(related_values).reshape(-1).tolist()
    if set(value) != set(flattened_vals):
        return get_related_values(data, flattened_vals) 
    else:
        return related_values

修改说明 #

1. 输入格式统一：通过isinstance判断输入类型，将int/float转为单元素列表，确保data.isin(value)能正常执行。
2. 递归参数优化：把展平后的结果数组转为list再传入递归，避免后续调用出现非list类型的输入，保持逻辑一致性。

测试示例 #

调用get_related_values(data, 129)，会先找到user_id=129的行，再通过该行的phone_number="61477708777"找到关联的第4行，最终返回结果：

id  user_id phone_number           email
1   2      129  61477708777   acdc@gmail.com
3   4      145  61477708777  austr@gmail.com

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