How to restrict DataFrame number of rows to the Xth unique value in certain column?

Question

Say for example we have the following DataFrame:

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

And we would know we wanted an x(say 3) number of unique values in column A. Then the desired output would be:

A B
1 2
1 2
2 3
3 4

I thought about looping through the column in question, counting the number of unique values by tracking and taking the subset of the DataFrame with the right index. I am still a newbie to Python and I believe there would be a more efficient way to do this, please share your solutions. Appreciated!

Answer 1

You can try series.factorize which indexes the unique values starting at 0 and then select the values which is <= n-1 ( because index starts at 0 ),hence reserves order too:

n=3
df[df['A'].factorize()[0]<=n-1]

---

   A  B
0  1  2
1  1  2
2  2  3
3  3  4

Answer 2

You can use np.random.choice to select the unique id, then isin to select rows with those id:

selected_ids = np.random.choice(df['A'].unique(), replace=False, size=3)

df[df['A'].isin(selected_ids)]