pandas dataframe how to remove values from cell that is a list based on other column-CodePudding

I have a dataframe with 2 columns that represent a list:

a. b.  vals.        locs
1. 2. [1,2,3,4,5].  [2,3]
5  1. [1,7,2,4,9].  [0,1]
8. 2. [1,9,4,7,8].  [3]

I want, for each row, exclude from the columns vals all the locations that are in locs. so I will get:

a. b.  vals.        locs.   new_vals
1. 2. [1,2,3,4,5].  [2,3].  [1,2,5]
5  1. [1,7,2,4,9].  [0,1].  [2,4,9]
8. 2. [1,9,4,7,8].  [3].    [1,9,4,8]

What is the best way to do so?

Thanks!

CodePudding user response：

You can use a list comprehension with an internal filter based on enumerate:

df['new_vals'] = [[v for i,v in enumerate(a) if i not in b]
                  for a,b in zip(df['vals'], df['locs'])]

however this will become quickly inefficient when b get large.

A much better approach would be to use python sets that enable a fast (O(1) complexity) identification of membership:

df['new_vals'] = [[v for i,v in enumerate(a) if i not in S]
                  for a,b in zip(df['vals'], df['locs']) for S in [set(b)]]

output:

   a  b             vals    locs      new_vals
0  1  2  [1, 2, 3, 4, 5]  [2, 3]     [1, 2, 5]
1  5  1  [1, 7, 2, 4, 9]  [0, 1]     [2, 4, 9]
2  8  2  [1, 9, 4, 7, 8]     [3]  [1, 9, 4, 8]

CodePudding user response：

Use list comprehension with enumerate and converting values to sets:

df['new_vals'] = [[z for i, z in enumerate(x) if i not in y]
                              for x, y in zip(df['vals'], df['locs'].apply(set))]
print (df)
   a  b             vals    locs      new_vals
0  1  2  [1, 2, 3, 4, 5]  [2, 3]     [1, 2, 5]
1  5  1  [1, 7, 2, 4, 9]  [0, 1]     [2, 4, 9]
2  8  2  [1, 9, 4, 7, 8]     [3]  [1, 9, 4, 8]

CodePudding user response：

One way to do this is to create a function that works on row,

def func(row):
    ans = [v for v in row['vals'] if row['vals'].index(v) not in row['locs']]
    return ans

The call this function for each row using apply.

df['new_value'] = df.apply(func, axis=1)

This will work well, if the lists are short.