Python中基于DataFrame列关系的索引匹配逻辑实现求助

原始数据 #

给定如下升序排列的DataFrame：

需求逻辑 #

需要按以下规则生成结果列：

• 若A < B，返回1.0
• 若A = B，返回对应行的value值
• 若A > B，找到B列中最接近A的两个较小值，返回这两个值的下一个B值对应的value

解决方案 #

利用numpy的np.where实现条件分支，结合np.searchsorted处理第三部分的核心逻辑（因B列是升序排列，np.searchsorted可快速定位第一个大于A的B值的索引，该索引对应的value即为需求结果）。

完整代码如下：

import pandas as pd
import numpy as np

# 创建示例DataFrame
data = {
    'value': [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
    'A': [7.0, 9.0, 9.5, 10.0, 13.0, 15.0, 16.0, 17.0],
    'B': [8.0, 8.8, 9.1, 9.4, 9.7, 9.9, 10.6, 17.0]
}
df = pd.DataFrame(data)

# 定义三个条件分支
cond_a_lt_b = df['A'] < df['B']
cond_a_eq_b = df['A'] == df['B']
cond_a_gt_b = df['A'] > df['B']

# 定位第一个大于A的B值的索引，处理极端情况（A大于所有B时取最后一个索引）
pos = np.searchsorted(df['B'], df['A'])
pos = np.clip(pos, 0, len(df)-1)

# 嵌套np.where实现逻辑
result = np.where(
    cond_a_lt_b,
    1.0,
    np.where(
        cond_a_eq_b,
        df['value'],
        df['value'].iloc[pos].values
    )
)

# 将结果添加到DataFrame
df['result'] = result
print(df)

验证结果 #

运行代码后得到的结果如下，完全符合需求：

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