Pandas: Keep values in a set of columns if they exist in another set of columns in the same row, otherwise set it to NaN

user_5285822 0
Viewed 177

I have a couple of columns in my dataframe that have values in them. I want to only keep those values in those columns if they exist in another set of columns in the same row. Otherwise, I want to set the value to NaN.

Here's an example dataframe:

    A   B   C   D
0   1  30   1  29
1   5  42  99   5
2  64  67  12  22
3   2  22  22   0
4  43   6   9  43

In this case, I want C and D to be changed based on A and B:

    A   B     C     D
0   1  30   1.0   NaN
1   5  42   NaN   5.0
2  64  67   NaN   NaN
3   2  22  22.0   NaN
4  43   6   NaN  43.0

It's been difficult to form a query to google this, and the closest I've gotten is to use pandas.DataFrame.isin like this:

from operator import concat
first = df.head(1)
first[['C', 'D']].isin(reduce(concat, first[['A', 'B']].values.tolist()))

Which gives me this:

      C      D
0  True  False

Which appears to be somewhat useful, but I'm not sure if this is the right path or what to do with it from here.

pandas dataframe python
2 Answers

Numpy broadcasting and pd.DataFrame.where

cd = df[['C', 'D']].to_numpy()
ab = df[['A', 'B']].to_numpy()

df[['C', 'D']] = df[['C', 'D']].where((cd[..., None] == ab[:, None]).any(axis=2))

df

    A   B     C     D
0   1  30   1.0   NaN
1   5  42   NaN   5.0
2  64  67   NaN   NaN
3   2  22  22.0   NaN
4  43   6   NaN  43.0

Less Numpy

df[['C', 'D']] = [
    (c if c in ab else np.nan, d if d in ab else np.nan)
    for *ab, c, d in zip(*map(df.get, df))
]

df

    A   B     C     D
0   1  30   1.0   NaN
1   5  42   NaN   5.0
2  64  67   NaN   NaN
3   2  22  22.0   NaN
4  43   6   NaN  43.0

Same thing but more specific with the columns

df[['C', 'D']] = [
    (c if c in ab else np.nan, d if d in ab else np.nan)
    for *ab, c, d in zip(*map(df.get, ['A', 'B', 'C', 'D']))
]
user_23366540

You need two masks, with A and B and OR between them. 

m1 = df[['C', 'D']] == pd.DataFrame({'C':df['A'], 'D':df['A']})
m2 = df[['C', 'D']] == pd.DataFrame({'C':df['B'], 'D':df['B']})
df[['C', 'D']] = df[['C', 'D']][(m1 | m2)]

Output:

    A   B     C     D
0   1  30   1.0   NaN
1   5  42   NaN   5.0
2  64  67   NaN   NaN
3   2  22  22.0   NaN
4  43   6   NaN  43.0
user_96094470
Related