Add list as a column to a dataframe

Viewed 145

Suppose I have the following lists:

cond_1 = [1,2]
cond_2 = [3,5]

And also the following dataframe df:

|----------|
| Column_1 |
|----------|
|     x    |
|----------|
|     y    |
|----------|
|     y    |
|----------|
|     x    |
|----------|

What I want to do, is to add a second column Column_2. following these criteria:

1) if Column_1 contains a x, add a value in Column_2 from cond_1;

2) if Column_1 contains a y, add a value in Column_2 from cond_2

The desired output should be like this:

|----------|----------|
| Column_1 | Column_2 |
|----------|----------|
|     x    |     1    |
|----------|----------|
|     y    |     3    |
|----------|----------|
|     y    |     5    |
|----------|----------|
|     x    |     2    |
|----------|----------|

I have been trying to do this using pd.Series:

df_x = df.loc[df['Column_1'] == "x"] #first I create a dataframe only with the x values
df_x['Column_2'] = pd.Series(cond_1)

Then I would repeat the same thing for the y values, obtaining df_y.

However, this doesn't succeed. Then, I would need to append again the two dataframes (df_x and df_y) and I lose information on the original index that I want to maintain from df.

4 Answers

You can create a helper class and use it in an .apply, eg:

class ReplaceWithNext:
    def __init__(self, **kwargs):
        self.lookup = {k: iter(v) for k, v in kwargs.items()}
    def __call__(self, value):
        return next(self.lookup[value])

Then use it as:

df['Column_2' ] = df['Column_1'].apply(ReplaceWithNext(x=cond_1, y=cond_2))

Which'll give you:

  Column_1  Column_2
0        x         1
1        y         3
2        y         5
3        x         2

A solution with loop :

choice = ['x','y']
cond_1 = [1,2]
cond_2 = [3,5]
d = dict(zip(choice,np.vstack((cond_1,cond_2))))
#{'x': array([1, 2]), 'y': array([3, 5])}

for k,v in d.items():
    df.loc[df['Column_1'].eq(k),'Column2'] = v
print(df)

  Column_1  Column2
0        x      1.0
1        y      3.0
2        y      5.0
3        x      2.0

You can merge. pd.concat will properly enumerate the index of each element of the lists. We'll need to groupby + cumcount the DataFrame to create this key there.

s = pd.concat([pd.Series(l).rename('Column_2') for l in [cond_1, cond_2]], 
              keys=['x', 'y'], names=['Column_1', 'N'])

df['N'] = df.groupby('Column_1').cumcount()
df = df.merge(s, on=['Column_1', 'N'], how='left').drop(columns='N')

  Column_1  Column_2
0        x         1
1        y         3
2        y         5
3        x         2

Using the keys and names arguments of pd.concat we can set up the merging Series which looks like below

print(s)
Column_1  N
x         0    1
          1    2
y         0    3
          1    5
Name: Column_2, dtype: int64
df = pd.DataFrame({'Column_1':['x', 'y', 'y', 'x'], 'Column_2':['','','','']})

cond_1 = [1,2]
cond_2 = [3,5]

cond_1_idx = 0
cond_2_idx = 0

col_2_list = []
for idx, row in df.iterrows():
    if df.at[idx ,'Column_1'] == 'x':

        col_2_list.append(cond_1[cond_1_idx])
        cond_1_idx +=1

    if df.at[idx ,'Column_1'] == 'y':

        print( df.at[0 ,'Column_1'])
        col_2_list.append(cond_2[cond_2_idx])
        cond_2_idx +=1

df['Column_2'] = col_2_list

    Column_1    Column_2
0   x   1
1   y   3
2   y   5
3   x   2
Related