如何在其他列上迭代矢量化 if/else 语句？

Posted 2023-03-11

【中文标题】如何在其他列上迭代矢量化 if/else 语句？

【英文标题】：How to iterate a vectorized if/else statement over additional columns?

【发布时间】：2018-10-16 18:09:56

【问题描述】：

import pandas as pd, numpy as np

ltlist = [1, 2]
org = 'ID': [1, 3, 4, 5, 6, 7], 'ID2': [3, 4, 5, 6, 7, 2]

ltlist_set = set(ltlist)
org['LT'] = np.where(org['ID'].isin(ltlist_set), org['ID'], 0)

我需要检查 ID2 列并将 ID 写入其中，除非它已经有 ID。

输出

ID  ID2 LT
1   3   1
3   4   0
4   5   0
5   6   0
6   7   0
7   2   2

谢谢！

【参考方案1】：

选项 1

您可以嵌套numpy.where 语句：

org['LT'] = np.where(org['ID'].isin(ltlist_set), 1,
                     np.where(org['ID2'].isin(ltlist_set), 2, 0))

选项 2

或者，您可以按顺序使用pd.DataFrame.loc：

org['LT'] = 0  # default value
org.loc[org['ID2'].isin(ltlist_set), 'LT'] = 2
org.loc[org['ID'].isin(ltlist_set), 'LT'] = 1

选项 3

第三种选择是使用numpy.select：

conditions = [org['ID'].isin(ltlist_set), org['ID2'].isin(ltlist_set)]
values = [1, 2]

org['LT'] = np.select(conditions, values, 0)  # 0 is default value