Insert list into cells which meet column conditions

Question

Consider df

   A  B  C
0  3  2  1
1  4  2  3
2  1  4  1
3  2  2  3

I want to add another column "D" such that D contains different Lists based on conditions on "A", "B" and "C"

   A  B  C  D
0  3  2  1  [1,0]
1  4  2  3  [1,0]
2  1  4  1  [0,2]
3  2  2  3  [2,0]

My code snippet looks like:

df['D'] = 0
df['D'] = df['D'].astype(object)

df.loc[(df['A'] > 1) & (df['B'] > 1), "D"] = [1,0]
df.loc[(df['A'] == 1) , "D"] = [0,2]
df.loc[(df['A'] == 2) & (df['C'] != 0) , "D"] = [2,0]

When I try to run this code it throws the following error:

ValueError: Must have equal len keys and value when setting with an iterable

I have converted the column into Object type as suggested here but still with error.

What I can infer is that pandas is trying to iterate over the elements of the list and assigns each of those values to the cells where as I am trying to assign the entire list to all the cells meeting the criterion.

Is there any way I can assign lists in the above fashion?

Answer 1

Another solution is create Series filled by list with shape for generating length of df:

df.loc[(df['A'] > 1) & (df['B'] > 1), "D"] = pd.Series([[1,0]]*df.shape[0])
df.loc[(df['A'] == 1) , "D"] = pd.Series([[0,2]]*df.shape[0])
df.loc[(df['A'] == 2) & (df['C'] != 0) , "D"] = pd.Series([[2,0]]*df.shape[0])
print (df)
   A  B  C       D
0  3  2  1  [1, 0]
1  4  2  3  [1, 0]
2  1  4  1  [0, 2]
3  2  2  3  [2, 0]

Answer 2

Here's a goofy way to do it

cond1 = df.A.gt(1) & df.B.gt(1)
cond2 = df.A.eq(1)
cond3 = df.A.eq(2) & df.C.ne(0)

df['D'] = cond3.map({True: [2, 0]}) \
  .combine_first(cond2.map({True: [0, 2]})) \
  .combine_first(cond1.map({True: [1, 0]})) \

df