Как в DataFrame сравнить два столбца с np.nan?

Question

Есть DataFrame, в нем надо сравнить два столбца так, чтобы сравнение np.nan и np.nan в двух столбцах выдавало True.

Код ниже возвращает две строки.
np.nan сравнивается как False.

df = pd.DataFrame({'Город1': ['Москва', 'Ярославль', 'Псков', np.nan,1],
    'Город2': ['Москва', 'Самара', 'Ростов', np.nan,1]})
df[ df['Город1'] == df['Город2'] ]

Написать сравнение следующим образом мне не подходит:

df[ (df['Город1'] == df['Город2']) | (df['Город1'].isnull()& df['Город2'].isnull()) ]

Я хочу сделать сравнение одним условием.

Спасибо!

Answer 1

Вероятно, самый простой способ - с fillna:

res = df[df["Город1"].fillna("nan")==df["Город2"].fillna("nan")].replace("nan", np.nan)

res:

   Город1  Город2
0  Москва  Москва
3     NaN     NaN
4       1       1

Answer 2

Предлагаю рассмотреть вариант с использованием pandas.DataFrame.nqunique У этого метода есть параметр dropna, который по умолчанию равен True. Нужно применить метод с dropna=False и отобрать те строки, где результат равен единице (значения nan в разных столбцах будут посчитаны как одно уникальное):

df[df[['Город1', 'Город2']].nunique('columns', dropna=False).eq(1)]

Код для экспериментов:

import pandas as pd

NA = float('nan')
df = pd.DataFrame({
    'Город1': ['Москва', 'Ярославль',  'Псков', NA, 1, NA, 'Иркутск'],
    'Город2': ['Москва',    'Самара', 'Ростов', NA, 1, NA, 'Иркутск'],
    'Город3': ['Москва',          NA, 'Ростов', NA, 1,  1,        NA],
})

print(df)

#       Город1  Город2  Город3
# 0     Москва  Москва  Москва
# 1  Ярославль  Самара     NaN
# 2      Псков  Ростов  Ростов
# 3        NaN     NaN     NaN
# 4          1       1       1
# 5        NaN     NaN       1
# 6    Иркутск Иркутск     NaN

cols_compare = 'Город2 Город3'.split()
selected = df[cols_compare].nunique(1, False).eq(1)
print(df.loc[selected, cols_compare])

#    Город2  Город3
# 0  Москва  Москва
# 2  Ростов  Ростов
# 3     NaN     NaN
# 4       1       1

print(df[df.nunique(1, False).eq(1)])

#    Город1  Город2  Город3
# 0  Москва  Москва  Москва
# 3     NaN     NaN     NaN
# 4       1       1       1

БЛОГ НА HUSL

Как в DataFrame сравнить два столбца с np.nan?

Ответы (2 шт):