Как перебрать векторизованный оператор if / else для дополнительных столбцов?

Question

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  ID2 LT
1   3   1
3   4   0
4   5   0
5   6   0
6   7   0
7   2   2

Спасибо!

Answer 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