选择具有两个唯一标识符的 pandas dataframe 中的行并将它们存储为新数据帧

Question

I have an extremely large, unsorted pandas dataframe (over two million rows) with multiple columns, two columns of which identify which category these rows belong to. Where the combination of "K" and "U" represent a unique category for these rows, I want to select all the rows that fall into each of these categories, and store these rows as separate dataframes that can be manipulated and analyzed later on for machine learning models. Let me explain

'a' 'b' 'c' 'K' 'U' 'd'
------------------------
aaa bbb ccc 2245 23 ddd
avd bad cec 2245 23 dwq
avd bad cec 2646 23 dwq
avd bad cec 1621 23 dwq
avd bad cec 1621 26 dwq

The two uppermost rows have the same "K" and "U" value, so I want these to be stored together, however, the other rows all belong to different categories altogether due to having a different combination of K and U , so these will be stored in a separate dataframe.

My first "solution" for this used a for loop to iterate through the dataframe's K , making a new dataframe encapsulating every row that contains the unique K , and making another for loop for every U in this new K dataframe. I then created a second dataframe in this loop containing that contains every row with this current U . This approach does not work as intended, but I feel I was close to a solution to the problem. It is unbearably slow on the full dataframe, and a quicker, proper solution would be appreciated. How would I go about doing this in a proper, more efficient manner?

Answer 1

You can do it this way:

         c     K   U    d
     0  aaa  2245  23  ddd
     1  avd  2245  23  dwq
     2  avd  2646  23  dwq
     3  avd  1621  23  dwq
     4  avd  1621  26  dwq 

grouped_df = dataframe.groupby(['K','U'])
for key,df in grouped_df:
  print('\n',key,'\n',df.head())

(1621, 23) 
      c     K   U    d
  3  avd  1621  23  dwq

(1621, 26) 
      c     K   U    d
  4  avd  1621  26  dwq

(2245, 23) 
      c     K   U    d
  0  aaa  2245  23  ddd
  1  avd  2245  23  dwq

(2646, 23) 
      c     K   U    d
  2  avd  2646  23  dwq

In this way you have n different dataframes with same pair value of 'K' and 'U'. After grouping you can access the single dataframe using the get_group method providing the key like:

df_n=grouped_df.get_group((2245, 23))

(2245, 23) 
      c     K   U    d
  0  aaa  2245  23  ddd
  1  avd  2245  23  dwq

Answer 2

You can use duplicated , specify a subset of the key columns and pass keep=False . Then put all of this inside of df[] to filter for those rows:

df[df.duplicated(subset=['K', 'U'], keep=False)]

    a   b   c   K       U   d
0   aaa bbb ccc 2245    23  ddd
1   avd bad cec 2245    23  dwq

For the other dataframe just add a ~ in front:

df[~df.duplicated(subset=['K', 'U'], keep=False)]

    a   b   c   K       U   d
2   avd bad cec 2646    23  dwq
3   avd bad cec 1621    23  dwq
4   avd bad cec 1621    26  dwq