Populate a pandas column with a list of values based on condition

May 2, 2023

I am trying populate a column with a list of values based on a condition. For example:

U=[2,3]
S=[1,3]
R=[1,2]
test = pd.DataFrame({'a': [1, 2, 3]})

test.loc[test.a.isin([1]), 'U']=U
test.loc[test.a.isin([2]), 'U']=S
test.loc[test.a.isin([3]), 'U']=R

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

I realize it’s the multiple values in the lists throwing the error. I’m guessing I should use lambda but wasn’t sure how to go about. Any suggestions. Thanks.

>Solution :

Since an example desired output isn’t provided, I’m not totally clear what you’re asking for.

But it sounds like you’re trying to map some values to another? So starting from your example:

import pandas as pd

U=[2,3]
S=[1,3]
R=[1,2]

test = pd.DataFrame({'a': [1, 2, 3]})

mapping = dict(zip(test.a, [U, S, R]))
test['a'].map(mapping)
# 0    [2, 3]
# 1    [1, 3]
# 2    [1, 2]
# Name: a, dtype: object

test['U'] = test['a'].map(mapping)
test
#    a       U
# 0  1  [2, 3]
# 1  2  [1, 3]
# 2  3  [1, 2]