2

After looking here and here and in the documentation, I still cannot find a way to select rows from a DataFrame according to all these criteria:

  • Return rows in an order given from a list of values from a given column
  • Return repeated rows (associated with repeated values in the list)
  • Preserve the original indices
  • Ignore values of the list not present in the DataFrame

As an example, let

df = pd.DataFrame({'A': [5, 6, 3, 4], 'B': [1, 2, 3, 5]})
df

     A   B
0    5   1
1    6   2
2    3   3
3    4   5

and let

list_of_values = [3, 4, 6, 4, 3, 8]

Then I would like to get the following DataFrame:

     A   B
2    3   3
3    4   5
1    6   2
3    4   5
2    3   3

How can I accomplish that? Zero's answer looks promising as it is the only one I found which preserves the original index, but it does not work with repetitions. Any ideas about how to modify/generalize it?

Lightspark
  • 355
  • 3
  • 10

2 Answers2

1

Here's a way to do that using merge:

list_df = pd.DataFrame({"A": list_of_values, "order": range(len(list_of_values))})

pd.merge(list_df, df, on="A").sort_values("order").drop("order", axis=1)

The output is:

   A  B
0  3  3
2  4  5
4  6  2
3  4  5
1  3  3
Roy2012
  • 11,755
  • 2
  • 22
  • 35
1

We have to preserve the index by assigning it as a column first so we can set_index after the mering:

list_of_values = [3, 4, 6, 4, 3, 8]
df2 = pd.DataFrame({'A': list_of_values, 'order': range(len(list_of_values))})

dfn = (
    df.assign(idx=df.index)
    .merge(df2, on='A')
    .sort_values('order')
    .set_index('idx')
    .drop('order', axis=1)
)

     A  B
idx      
2    3  3
3    4  5
1    6  2
3    4  5
2    3  3

If you want to get rid of the index name (idx), use rename_axis:

dfn = dfn.rename_axis(None)

   A  B
2  3  3
3  4  5
1  6  2
3  4  5
2  3  3
Erfan
  • 40,971
  • 8
  • 66
  • 78
  • Thank you, this works indeed. As for renaming the index, I suppose that also `dfn.index = dfn.index.rename("")` would work, but it is surely more verbose. – Lightspark Jul 31 '20 at 14:27