CODEWITHSUNDEEP
pythonpython-3.xpandasdataframenumpy
William
William

Reputation: 4034

How add multiple columns with a condition Using np.where()

I know how to use np.where() to add one column by 1 condition:

import pandas as pd
import numpy as np
df=pd.read_csv(file,nrows=5)
df['new_col1']= np.where(df['col1'] < '100', 1,2)
df.head()

output:

   col1  col2  new_col1
0     1     3    1
1     2     4    1

what if I want to add 2 columns by the same condition:

df['new_col1'],df['new_col2']= np.where(df['col1'] < '100', (1,2),(3,4))

I want to add new_col1 and new_col2,the result are (1,2),(3,4)

When I tried this code, I received:

ValueError: too many values to unpack (expected 2)

The output should be:

   col1  col2  new_col1 new_col2
0     1     3    1       3
1     2     4    1       3

Upvotes: 0

Views: 703

Answers (1)

Andreas
Andreas

Reputation: 9207

You can use the condition multiple times:

mask = df['contract'] > '0L000099'
df['column1'] = np.where(mask, 1, 2)
df['column2'] = np.where(mask, 3, 4)

or even invert the condition:

df['column2'] = np.where(~mask, 1, 2)

Since your question was updated, here the updated answer, however I am not sure thats actually usefull:

import pandas as pd
df = pd.DataFrame({'test':range(0,10)})
mask  = df['test'] > 3
m_len = len(mask)

df['column1'], df['column2'] = np.where([mask, mask], [[1]*m_len, [3]*m_len], [[2]*m_len, [4]*m_len])

   test  column1  column2
0     0        2        4
1     1        2        4
2     2        2        4
3     3        2        4
4     4        1        3
5     5        1        3
6     6        1        3
7     7        1        3
8     8        1        3
9     9        1        3

Upvotes: 1

Related Questions