Home > Back-end >  Adding list item to DataFrame column if it exists within a range
Adding list item to DataFrame column if it exists within a range

Time:12-14

I'm attempting to add list items to specified columns within a DataFrame if a list item exists in the specified position.

list1 = ['a', 'b', 'c', 'd', 'e']

df = pd.DataFrame(columns = ['Letter', 'Headline 1', 'Headline 2', 'Headline 3', 'Headline 4', 'Headline 5', 'Headline 6'])
data = {'Letter':['one', 'two', 'three', 'four', 'five']}
df = pd.DataFrame(data)

I'm attempting to do so with a series of if statements. It works if there is a list item in the specified position. I'm sure that this isn't an ideal way to accomplish what I am looking for. It also doesn't work if one of the list items is out of range. Ideally, I would like to insert an empty string if the list item doesn't exist.

if list1[0]:
    df['Headline1'] = list1[0]
    
if list1[1]:
    df['Headline2'] = list1[1]
    
if list1[2]:
    df['Headline3'] = list1[2]
    
if list1[3]:
    df['Headline4'] = list1[3]
    
if list1[4]:
    df['Headline5'] = list1[4]
    
if list1[5]:
    df['Headline6'] = list1[5]

Desired Output:

  Letter Headline1 Headline2 Headline3 Headline4 Headline5 Headline6
0    one     a          b        c          d          e           
1    two     a          b        c          d          e           
2  three     a          b        c          d          e           
3   four     a          b        c          d          e           
4   five     a          b        c          d          e    

CodePudding user response:

Use DataFrame.assign with dictionary created by zip:

L = [f'Headline {x 1}' for x in range(list1)]
d = dict(zip(L, list1))
df = pd.DataFrame(data).assign(**d)
print (df)
  Letter Headline 1 Headline 2 Headline 3 Headline 4 Headline 5
0    one          a          b          c          d          e
1    two          a          b          c          d          e
2  three          a          b          c          d          e
3   four          a          b          c          d          e
4   five          a          b          c          d          e

L = ['Headline 1', 'Headline 2', 'Headline 3', 'Headline 4', 'Headline 5', 'Headline 6']
d = dict(zip(L, list1))
df = pd.DataFrame(data).assign(**d)
print (df)
  Letter Headline 1 Headline 2 Headline 3 Headline 4 Headline 5
0    one          a          b          c          d          e
1    two          a          b          c          d          e
2  three          a          b          c          d          e
3   four          a          b          c          d          e
4   five          a          b          c          d          e

CodePudding user response:

I'd use assign with unpacking and dropna:

df.assign(**data).fillna(dict(zip(df.columns[1:], list1))).dropna(axis=1)

Output:

  Letter Headline 1 Headline 2 Headline 3 Headline 4 Headline 5
0    one          a          b          c          d          e
1    two          a          b          c          d          e
2  three          a          b          c          d          e
3   four          a          b          c          d          e
4   five          a          b          c          d          e
  • Related