Filtering pandas dataframe lists using list of indexes

Question

I have a dataframe which each cell have a list of values. In another dictionnary, I have a list of indexes I want to keep for each row of my dataframe. I want to filter my dataframe using my dictionnary.

df = {"a": [[1,2,3,4], [4,5,6,7]],
      "b": [[11,22,33,44],[44,55,66,77]],
      "c": [[111,222,333,444],[444,555,666,777]]}

df = pd.DataFrame(df)

filtered_idx = {0: [0,2],
                1: [3],
                2: []}

filtered_df = df.applymap(func(filtered_idx))

expected output:

   a        b
0  [1,3]  [4,6]
1  [44]   [77]
2  []     []

How do i need to implement my func function.

Thanks

Answer 1

You can try apply on rows and then use another apply on each row value

filtered_df = df.apply(lambda row: row.apply(lambda x: [x[idx] for idx in filtered_idx[row.name]]), axis=1)\

# or

filtered_df = df.apply(lambda row: row.apply(lambda x: np.array(x)[filtered_idx[row.name]]), axis=1)

        a         b           c
0  [1, 3]  [11, 33]  [111, 333]
1     [7]      [77]       [777]