How to pop rows from dataframe and from lists within a column

user_7460197 0
Viewed 393

I have a dataframe with a list as one of the columns. The list is common to the label.

    df = pd.DataFrame({'label': ['cell_1', 'cell_1', 'cell_1', 'cell_1', 'cell_2', 'cell_2', 'cell_2'],
                   'index': [1, 2, 5, 6, 1, 18, 5],
                   'item': [121, 989, 222, 333, 878, 777, 878],
                   'list': [[121, 454, 989], [121, 454, 989], [121, 454, 989], [121, 454, 989], [111, 777, 999], [111, 777, 999], [111, 777, 999]]})

    label    index  item    list
0   cell_1   1      121     [121, 454, 989]
1   cell_1   2      989     [121, 454, 989]
2   cell_1   5      222     [121, 454, 989]
3   cell_1   6      333     [121, 454, 989]
4   cell_2   1      878     [111, 777, 999]
5   cell_2   18     777     [111, 777, 999]
6   cell_2   5      878     [111, 777, 999]

I would like to check if the value in the item column is in the list. If it is I would like to pop the row into its own dataframe:

popped_df

   label    index   item
0  cell_1   1       121
1  cell_1   2       989
2  cell_2   18      777

And remove the item value from the list in the remaining dataframe:

df
   label    index   item    list
0  cell_1   5       222     [454]
1  cell_1   6       333     [454]
2  cell_2   1       878     [111, 999]
3  cell_2   5       878     [111, 999]

I have tried many things, all of which run into big problems. I think maybe the solution is going to be putting the data in nested dictionaries (with label as the key, and each index as a value) and iterating through. I would rather not go from dataframe to dictionary and back even if it means slow performance.

Thank you in advance!

pandas dataframe list
3 Answers
  1. You can do this with list comprehension if you want to keep it in a list (rather than explode the dataframe). Since you are comparing multiple columns, use .apply(lambda x:) on the dataframe, where x is the dataframe and pass the columns and axis=1 to go column-wise over the default axis=0. The performance might be better with explode, but I would test this method versus an explode method.
  2. To get rid of the values for the initial dataframe use a very similar list comprehension technique
df = pd.DataFrame({'label': ['cell_1', 'cell_1', 'cell_1', 'cell_1', 'cell_2', 'cell_2', 'cell_2'],
                   'index': [1, 2, 5, 6, 1, 18, 5],
                   'item': [121, 989, 222, 333, 878, 777, 878],
                   'list': [[121, 454, 989], [121, 454, 989], [121, 454, 989], [121, 454, 989], [111, 777, 999], [111, 777, 999], [111, 777, 999]]})

s = df.apply(lambda x: any([True for y in x['list'] if x['item'] == y ]), axis=1)
df1 = df[s].drop('list', axis=1)
df['list'] = df.apply(lambda x: [y for y in x['list'] if y not in df1['item'].unique()], axis=1)
df = df[~s]
df1, df

Out[1]: 
(    label  index  item
 0  cell_1      1   121
 1  cell_1      2   989
 5  cell_2     18   777,
     label  index  item        list
 2  cell_1      5   222       [454]
 3  cell_1      6   333       [454]
 4  cell_2      1   878  [111, 999]
 6  cell_2      5   878  [111, 999])
user_63667700

I recommend you explode the list and then do a basic boolean indexer:


(
    df.explode('list')
        .loc[lambda df: ~df['list'].isin(df['item']) & ~df['item'].isin(df['list'])]
        .pivot_table(index=['label', 'index', 'item'], values='list', aggfunc=list)
        .reset_index()
)

Which gives me:

    label  index  item        list
0  cell_1      5   222       [454]
1  cell_1      6   333       [454]
2  cell_2      1   878  [111, 999]
3  cell_2      5   878  [111, 999]

But I recommend you keep your values out of the list (omit the call to pivot_table)

    label  index  item list
2  cell_1      5   222  454
3  cell_1      6   333  454
4  cell_2      1   878  111
4  cell_2      1   878  999
6  cell_2      5   878  111
6  cell_2      5   878  999

You can flip the logic to get your other dataframe:

(
    df.explode('list')
        .loc[lambda df: (df['list'].isin(df['item']) & df['item'].isin(df['list']))]
        .drop(columns=['list'])
        .drop_duplicates(subset=['label', 'index'])
        .to_string()
)

Which is:

    label  index  item
0  cell_1      1   121
1  cell_1      2   989
5  cell_2     18   777

Note that in general, so-called tidy data is much easier to work with:

df.explode('list').reset_index(drop=True)

     label  index  item list
0   cell_1      1   121  121
1   cell_1      1   121  454
2   cell_1      1   121  989
3   cell_1      2   989  121
4   cell_1      2   989  454
5   cell_1      2   989  989
6   cell_1      5   222  121
7   cell_1      5   222  454
8   cell_1      5   222  989
9   cell_1      6   333  121
10  cell_1      6   333  454
11  cell_1      6   333  989
12  cell_2      1   878  111
13  cell_2      1   878  777
14  cell_2      1   878  999
15  cell_2     18   777  111
16  cell_2     18   777  777
17  cell_2     18   777  999
18  cell_2      5   878  111
19  cell_2      5   878  777
20  cell_2      5   878  999
user_15527480

Paul and David's answers were very helpful, but unfortunately when I applied them to my full data set, I found that their solutions removed items from the entire list column, and not just the list common to the label. For example, if I add a list item that occurs in both the cell_1 and cell_2 lists:

df = pd.DataFrame({'label': ['cell_1', 'cell_1', 'cell_1', 'cell_1', 'cell_2', 'cell_2', 'cell_2'],
               'index': [1, 2, 5, 6, 1, 18, 5],
               'item': [121, 989, 222, 333, 878, 777, 878],
               'list': [[121, 454, 989], [121, 454, 989], [121, 454, 989], [121, 454, 989], [121, 777, 999], [121, 777, 999], [121, 777, 999]]})

label   index   item    list
0   cell_1  1   121 [121, 454, 989]
1   cell_1  2   989 [121, 454, 989]
2   cell_1  5   222 [121, 454, 989]
3   cell_1  6   333 [121, 454, 989]
4   cell_2  1   878 [121, 777, 999]
5   cell_2  18  777 [121, 777, 999]
6   cell_2  5   878 [121, 777, 999]

With both solutions, the list for cell_2 has "121" removed even though I only wanted it removed from the list cells common to cell_1:

s = df.apply(lambda x: any([True for y in x['list'] if x['item'] == y ]), axis=1)
df1 = df[s].drop('list', axis=1)
df['list'] = df.apply(lambda x: [y for y in x['list'] if y not in df1['item'].unique()], axis=1)
df = df[~s]
df1, df
(    label  index  item
 0  cell_1      1   121
 1  cell_1      2   989
 5  cell_2     18   777,
     label  index  item   list
 2  cell_1      5   222  [454]
 3  cell_1      6   333  [454]
 4  cell_2      1   878  [999]
 6  cell_2      5   878  [999])

What I ended up doing was exploding the data to make it tidy:

df.explode('list').reset_index(drop=True)
df
            label   index   item    list
0   cell_1  1       121     121
1   cell_1  1       121     454
2   cell_1  1       121     989
3   cell_1  2       989     121
4   cell_1  2       989     454
5   cell_1  2       989     989
6   cell_1  5       222     121
7   cell_1  5       222     454
8   cell_1  5       222     989
9   cell_1  6       333     121
10  cell_1  6       333     454
11  cell_1  6       333     989
12  cell_2  1       878     121
13  cell_2  1       878     777
14  cell_2  1       878     999
15  cell_2  18      777     121
16  cell_2  18      777     777
17  cell_2  18      777     999
18  cell_2  5       878     121
19  cell_2  5       878     777
20  cell_2  5       878     999

And then adding extra columns:

df['lookup1'] = df['label'] + "-" + df['item'].astype(str)
df.loc[(df['item'] == df['list']), 'lookup2'] = df['label'] + "-" + df['item'].astype(str)
df["selector"] = df['lookup1'].isin(df["lookup2"])

df

    label   index   item    list    lookup1      lookup2    selector
0   cell_1  1       121     121     cell_1-121  cell_1-121  True
1   cell_1  1       121     454     cell_1-121  NaN         True
2   cell_1  1       121     989     cell_1-121  NaN         True
3   cell_1  2       989     121     cell_1-989  NaN         True
4   cell_1  2       989     454     cell_1-989  NaN         True
5   cell_1  2       989     989     cell_1-989  cell_1-989  True
6   cell_1  5       222     121     cell_1-222  NaN         False
7   cell_1  5       222     454     cell_1-222  NaN         False
8   cell_1  5       222     989     cell_1-222  NaN         False
9   cell_1  6       333     121     cell_1-333  NaN         False
10  cell_1  6       333     454     cell_1-333  NaN         False
11  cell_1  6       333     989     cell_1-333  NaN         False
12  cell_2  1       878     121     cell_2-878  NaN         False
13  cell_2  1       878     777     cell_2-878  NaN         False
14  cell_2  1       878     999     cell_2-878  NaN         False
15  cell_2  18      777     121     cell_2-777  NaN         True
16  cell_2  18      777     777     cell_2-777  cell_2-777  True
17  cell_2  18      777     999     cell_2-777  NaN         True
18  cell_2  5       878     121     cell_2-878  NaN         False
19  cell_2  5       878     777     cell_2-878  NaN         False
20  cell_2  5       878     999     cell_2-878  NaN         False

At which point I can operate on the selector column. Surely there is a more elegant way to do this?

user_74601970
Related