Select-largest-n-records-by-group.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Select records from a DataFrame by the n largest values per group\n",
    "\n",
    "Based on answers to the following stackoverflow posts:\n",
    "- https://stackoverflow.com/questions/63431597/how-to-drop-certain-values-within-a-multi-level-index-python-pandas/63432079#63432079\n",
    "- https://stackoverflow.com/questions/27842613/pandas-groupby-sort-within-groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data is from the Stackoverflow question linked above\n",
    "\n",
    "data = {\n",
    "    'full_name':\n",
    "        ['Michael Thomas', 'Chris Godwin', 'DeAndre Hopkins', 'Julio Jones', 'Cooper Kupp', 'Julian Edelman', \n",
    "         'Kenny Golladay', 'Keenan Allen', 'Amari Cooper', 'D.J. Moore', 'Mike Evans', 'DeVante Parker', \n",
    "         'Jarvis Landry', 'Stefon Diggs', 'Tyler Lockett', 'Tyler Boyd', 'John Brown', 'Robert Woods', \n",
    "         'Courtland Sutton', 'Calvin Ridley', 'A.J. Brown', 'Terry McLaurin', 'Davante Adams', 'Cole Beasley', \n",
    "         'Michael Gallup', 'Tyreek Hill', 'Jamison Crowder', 'Larry Fitzgerald', 'Curtis Samuel', 'Deebo Samuel', \n",
    "         'Darius Slayton', 'Mike Williams', 'Mike Williams', 'Robby Anderson', 'Christian Kirk', \n",
    "         'Diontae Johnson', 'Chris Conley', 'Randall Cobb', 'Tyrell Williams', 'Marquise Brown', 'Sammy Watkins', \n",
    "         'Golden Tate', 'James Washington', 'Dede Westbrook', 'Danny Amendola', 'Sterling Shepard', \n",
    "         'Zach Pascal', 'Anthony Miller', 'Alshon Jeffery', 'Kenny Stills', 'T.Y. Hilton', 'Adam Thielen', \n",
    "         'Breshad Perriman', 'Mecole Hardman', 'JuJu Smith-Schuster', 'Hunter Renfrow', 'Brandin Cooks', \n",
    "         'Corey Davis', 'Auden Tate', 'Emmanuel Sanders', 'Emmanuel Sanders', 'Alex Erickson', 'Kendrick Bourne', \n",
    "         'Nelson Agholor', 'Preston Williams', 'Demarcus Robinson', 'Taylor Gabriel', 'Russell Gage', \n",
    "         'Adam Humphries', 'Allen Lazard', 'Allen Hurns', 'Demaryius Thomas', 'Marquez Valdes-Scantling', \n",
    "         'Albert Wilson', 'Bisi Johnson', 'Geronimo Allison', 'Paul Richardson', 'Keelan Cole', 'Cody Latimer', \n",
    "         'Jakobi Meyers', 'Josh Reynolds', 'Kelvin Harmon', 'Seth Roberts', 'Isaiah McKenzie', 'David Moore', \n",
    "         'Josh Gordon', 'Josh Gordon', \"Tre'Quan Smith\", 'Jarius Wright', 'Damiere Byrd', 'Keke Coutee', \n",
    "         'DaeSean Hamilton', 'Pharoh Cooper', 'Trey Quinn', 'Miles Boykin', 'Jaron Brown', 'Chester Rogers', \n",
    "         'Tavon Austin', 'KeeSean Johnson', 'Malik Turner', 'Bennie Fowler', 'Vyncint Smith', 'Jakeem Grant', \n",
    "         'Parris Campbell', 'Marvin Hall', 'Jake Kumerow', 'Justin Watson', 'Marquise Goodwin', 'Javon Wims', \n",
    "         'DeSean Jackson', 'Andy Isabella', 'Tim Patrick', 'Byron Pringle', 'Dante Pettis', 'Laquon Treadwell', \n",
    "         'J.J. Arcega-Whiteside', \"N'Keal Harry\", 'Justin Hardy', 'Kalif Raymond', 'Zay Jones', 'Zay Jones', \n",
    "         'Trevor Davis', 'Trevor Davis', 'Trevor Davis', 'Cordarrelle Patterson', 'Keelan Doss', 'Damion Ratley', \n",
    "         'Mack Hollins', 'Mack Hollins', 'Devin Smith', 'Dontrelle Inman', 'Dontrelle Inman', 'Christian Blake', \n",
    "         'Duke Williams', 'Antonio Callaway', 'Olamide Zaccheaus', 'Rashard Higgins', 'Braxton Berrios', \n",
    "         'DeAndre Carter', 'Robert Foster', 'Deon Cain', 'Deon Cain', 'Trent Sherfield', 'Andre Patton', \n",
    "         'Chris Hogan', 'Ryan Switzer', 'Diontae Spencer', 'Deonte Harris', 'Chad Beebe', 'Marcell Ateman', \n",
    "         'Travis Benjamin', 'KhaDarel Hodge', 'Ventell Bryant', 'Geremy Davis', 'Jason Moore', 'Devin Funchess', \n",
    "         'Johnny Holton', 'Cody Core', 'Donte Moncrief', 'Andre Roberts', 'Russell Shepard', 'Gunner Olszewski', \n",
    "         'Ryan Grant', 'C.J. Board', 'Chris Moore', 'Marqise Lee', 'Stanley Morgan', 'Riley Ridley', \n",
    "         'Fred Brown', 'DeAndrew White', 'Josh Bellamy', 'Ashton Dulin', 'Michael Walker', 'Mike Thomas', \n",
    "         'Brandon Zylstra', 'Austin Carr', 'Dwayne Harris', 'Krishawn Hogan', 'Quincy Enunwa', 'Greg Dortch', \n",
    "         'JoJo Natson', 'Juwann Winfree', 'Matthew Slater', 'Taywan Taylor'],\n",
    "    'rec_yards':\n",
    "        [1688, 1333, 1165, 1316, 1062, 1091, 1118, 1117, 1097, 1175, 1157, 1065, 1092, 1130, 1006, 987, 1060, \n",
    "         1067, 1060, 866, 927, 919, 904, 778, 1009, 799, 767, 759, 614, 700, 690, 963, 963, 761, 649, 626, 737, \n",
    "         747, 651, 569, 665, 608, 735, 588, 662, 537, 597, 651, 490, 561, 429, 418, 511, 508, 546, 503, 543, \n",
    "         557, 575, 477, 367, 513, 358, 363, 428, 425, 353, 378, 374, 408, 416, 433, 433, 292, 260, 270, 245, \n",
    "         294, 288, 359, 326, 332, 271, 247, 271, 287, 139, 178, 286, 285, 247, 232, 219, 198, 198, 220, 179, \n",
    "         176, 187, 245, 193, 189, 164, 127, 261, 212, 132, 186, 163, 159, 189, 204, 170, 109, 184, 169, 76, \n",
    "         155, 170, 126, 69, 83, 28, 0, 83, 133, 136, 125, 0, 113, 132, 43, 91, 58, 89, 93, 55, 104, 97, 64, 72, \n",
    "         52, 80, 56, 53, 27, 31, 24, 70, 70, 30, 57, 15, 38, 43, 32, 21, 28, 18, 20, 25, 34, 14, 31, 21, 18, \n",
    "         18, 15, 21, 20, 20, 17, 15, 14, 10, 9, 7, 4, -4, 0, 0, 0, 0, 0],\n",
    "    'team':\n",
    "        ['NO', 'TB', 'HOU', 'ATL', 'LA', 'NE', 'DET', 'LAC', 'DAL', 'CAR', 'TB', 'MIA', 'CLE', 'MIN', 'SEA', \n",
    "         'CIN', 'BUF', 'LA', 'DEN', 'ATL', 'TEN', 'WAS', 'GB', 'BUF', 'DAL', 'KC', 'NYJ', 'ARI', 'CAR', 'SF', \n",
    "         'NYG', 'LAC', 'LAC', 'NYJ', 'ARI', 'PIT', 'JAC', 'DAL', 'OAK', 'BAL', 'KC', 'NYG', 'PIT', 'JAC', \n",
    "         'DET', 'NYG', 'IND', 'CHI', 'PHI', 'HOU', 'IND', 'MIN', 'TB', 'KC', 'PIT', 'OAK', 'LA', 'TEN', 'CIN', \n",
    "         'SF', 'DEN', 'CIN', 'SF', 'PHI', 'MIA', 'KC', 'CHI', 'ATL', 'TEN', 'GB', 'MIA', 'NYJ', 'GB', 'MIA', \n",
    "         'MIN', 'GB', 'WAS', 'JAC', 'NYG', 'NE', 'LA', 'WAS', 'BAL', 'BUF', 'SEA', 'NE', 'SEA', 'NO', 'CAR', \n",
    "         'ARI', 'HOU', 'DEN', 'ARI', 'WAS', 'BAL', 'SEA', 'IND', 'DAL', 'ARI', 'SEA', 'NYG', 'NYJ', 'MIA', \n",
    "         'IND', 'DET', 'GB', 'TB', 'SF', 'CHI', 'PHI', 'ARI', 'DEN', 'KC', 'SF', 'MIN', 'PHI', 'NE', 'ATL', \n",
    "         'TEN', 'OAK', 'BUF', 'OAK', 'GB', 'MIA', 'CHI', 'OAK', 'CLE', 'PHI', 'MIA', 'DAL', 'LAC', 'IND', 'ATL', \n",
    "         'BUF', 'CLE', 'ATL', 'CLE', 'NYJ', 'HOU', 'BUF', 'PIT', 'IND', 'ARI', 'LAC', 'CAR', 'PIT', 'DEN', 'NO', \n",
    "         'MIN', 'OAK', 'LAC', 'CLE', 'DAL', 'LAC', 'LAC', 'IND', 'PIT', 'NYG', 'PIT', 'BUF', 'NYG', 'NE', 'OAK', \n",
    "         'JAC', 'BAL', 'JAC', 'CIN', 'CHI', 'DEN', 'CAR', 'NYJ', 'IND', 'JAC', 'LA', 'CAR', 'NO', 'OAK', 'NO', \n",
    "         'NYJ', 'CAR', 'LA', 'DEN', 'NE', 'CLE']\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>full_name</th>\n",
       "      <th>rec_yards</th>\n",
       "      <th>team</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Michael Thomas</td>\n",
       "      <td>1688</td>\n",
       "      <td>NO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Chris Godwin</td>\n",
       "      <td>1333</td>\n",
       "      <td>TB</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>DeAndre Hopkins</td>\n",
       "      <td>1165</td>\n",
       "      <td>HOU</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Julio Jones</td>\n",
       "      <td>1316</td>\n",
       "      <td>ATL</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Cooper Kupp</td>\n",
       "      <td>1062</td>\n",
       "      <td>LA</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         full_name  rec_yards team\n",
       "0   Michael Thomas       1688   NO\n",
       "1     Chris Godwin       1333   TB\n",
       "2  DeAndre Hopkins       1165  HOU\n",
       "3      Julio Jones       1316  ATL\n",
       "4      Cooper Kupp       1062   LA"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(data)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a new dataframe containing only the top four entries per team by highest rec_yards"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 1 - make a count of the items per group then select the top 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>rec_yards</th>\n",
       "      <th>full_name</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>team</th>\n",
       "      <th>count</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">WAS</th>\n",
       "      <th>1</th>\n",
       "      <td>919</td>\n",
       "      <td>Terry McLaurin</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>332</td>\n",
       "      <td>Kelvin Harmon</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>245</td>\n",
       "      <td>Paul Richardson</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>198</td>\n",
       "      <td>Trey Quinn</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TEN</th>\n",
       "      <th>1</th>\n",
       "      <td>927</td>\n",
       "      <td>A.J. Brown</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ATL</th>\n",
       "      <th>4</th>\n",
       "      <td>155</td>\n",
       "      <td>Justin Hardy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">ARI</th>\n",
       "      <th>1</th>\n",
       "      <td>759</td>\n",
       "      <td>Larry Fitzgerald</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>649</td>\n",
       "      <td>Christian Kirk</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>285</td>\n",
       "      <td>Damiere Byrd</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>219</td>\n",
       "      <td>Pharoh Cooper</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>127 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            rec_yards         full_name\n",
       "team count                             \n",
       "WAS  1            919    Terry McLaurin\n",
       "     2            332     Kelvin Harmon\n",
       "     3            245   Paul Richardson\n",
       "     4            198        Trey Quinn\n",
       "TEN  1            927        A.J. Brown\n",
       "...               ...               ...\n",
       "ATL  4            155      Justin Hardy\n",
       "ARI  1            759  Larry Fitzgerald\n",
       "     2            649    Christian Kirk\n",
       "     3            285      Damiere Byrd\n",
       "     4            219     Pharoh Cooper\n",
       "\n",
       "[127 rows x 2 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Make a sorted multi-index\n",
    "df_sorted = df.set_index(['team','rec_yards']).sort_index(ascending=False)\n",
    "\n",
    "# Add a dummy column containing all 1s\n",
    "df_sorted['count'] = 1\n",
    "\n",
    "# Turn it into a ranking by team\n",
    "df_sorted['count'] = df_sorted.groupby('team')['count'].cumsum()       \n",
    "\n",
    "# Only keep the first n rows per team\n",
    "n = 4\n",
    "df_sorted = df_sorted.loc[df_sorted['count'] <= n]\n",
    "\n",
    "# Re-arrange it however you prefer\n",
    "df_sorted.reset_index().set_index(['team','count'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 2 - Sort first, then use head with groupby"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>full_name</th>\n",
       "      <th>rec_yards</th>\n",
       "      <th>team</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>Terry McLaurin</td>\n",
       "      <td>919</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>81</th>\n",
       "      <td>Kelvin Harmon</td>\n",
       "      <td>332</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76</th>\n",
       "      <td>Paul Richardson</td>\n",
       "      <td>245</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>93</th>\n",
       "      <td>Trey Quinn</td>\n",
       "      <td>198</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>A.J. Brown</td>\n",
       "      <td>927</td>\n",
       "      <td>TEN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>117</th>\n",
       "      <td>Justin Hardy</td>\n",
       "      <td>155</td>\n",
       "      <td>ATL</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>Larry Fitzgerald</td>\n",
       "      <td>759</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>Christian Kirk</td>\n",
       "      <td>649</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>89</th>\n",
       "      <td>Damiere Byrd</td>\n",
       "      <td>285</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>92</th>\n",
       "      <td>Pharoh Cooper</td>\n",
       "      <td>219</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>127 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            full_name  rec_yards team\n",
       "21     Terry McLaurin        919  WAS\n",
       "81      Kelvin Harmon        332  WAS\n",
       "76    Paul Richardson        245  WAS\n",
       "93         Trey Quinn        198  WAS\n",
       "20         A.J. Brown        927  TEN\n",
       "..                ...        ...  ...\n",
       "117      Justin Hardy        155  ATL\n",
       "27   Larry Fitzgerald        759  ARI\n",
       "34     Christian Kirk        649  ARI\n",
       "89       Damiere Byrd        285  ARI\n",
       "92      Pharoh Cooper        219  ARI\n",
       "\n",
       "[127 rows x 3 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n = 4\n",
    "df.sort_values(['team','rec_yards'], ascending=False).groupby('team').head(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 3 - using `DataFrame.nlargest`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>full_name</th>\n",
       "      <th>rec_yards</th>\n",
       "      <th>team</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>team</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">ARI</th>\n",
       "      <th>27</th>\n",
       "      <td>Larry Fitzgerald</td>\n",
       "      <td>759</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>Christian Kirk</td>\n",
       "      <td>649</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>89</th>\n",
       "      <td>Damiere Byrd</td>\n",
       "      <td>285</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>92</th>\n",
       "      <td>Pharoh Cooper</td>\n",
       "      <td>219</td>\n",
       "      <td>ARI</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ATL</th>\n",
       "      <th>3</th>\n",
       "      <td>Julio Jones</td>\n",
       "      <td>1316</td>\n",
       "      <td>ATL</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TEN</th>\n",
       "      <th>118</th>\n",
       "      <td>Kalif Raymond</td>\n",
       "      <td>170</td>\n",
       "      <td>TEN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">WAS</th>\n",
       "      <th>21</th>\n",
       "      <td>Terry McLaurin</td>\n",
       "      <td>919</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>81</th>\n",
       "      <td>Kelvin Harmon</td>\n",
       "      <td>332</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>76</th>\n",
       "      <td>Paul Richardson</td>\n",
       "      <td>245</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>93</th>\n",
       "      <td>Trey Quinn</td>\n",
       "      <td>198</td>\n",
       "      <td>WAS</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>127 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                 full_name  rec_yards team\n",
       "team                                      \n",
       "ARI  27   Larry Fitzgerald        759  ARI\n",
       "     34     Christian Kirk        649  ARI\n",
       "     89       Damiere Byrd        285  ARI\n",
       "     92      Pharoh Cooper        219  ARI\n",
       "ATL  3         Julio Jones       1316  ATL\n",
       "...                    ...        ...  ...\n",
       "TEN  118     Kalif Raymond        170  TEN\n",
       "WAS  21     Terry McLaurin        919  WAS\n",
       "     81      Kelvin Harmon        332  WAS\n",
       "     76    Paul Richardson        245  WAS\n",
       "     93         Trey Quinn        198  WAS\n",
       "\n",
       "[127 rows x 3 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.groupby('team').apply(pd.DataFrame.nlargest, 4, columns=['rec_yards'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 4 - using `pd.SeriesGroupBy.nlargest`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "team     \n",
       "ARI   27      759\n",
       "      34      649\n",
       "      89      285\n",
       "      92      219\n",
       "ATL   3      1316\n",
       "             ... \n",
       "TEN   118     170\n",
       "WAS   21      919\n",
       "      81      332\n",
       "      76      245\n",
       "      93      198\n",
       "Name: rec_yards, Length: 127, dtype: int64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_n_by_team = df.groupby('team')['rec_yards'].nlargest(4)  # Only works for one series\n",
    "top_n_by_team"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>full_name</th>\n",
       "      <th>rec_yards</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>team</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>WAS</th>\n",
       "      <td>Terry McLaurin</td>\n",
       "      <td>919</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAS</th>\n",
       "      <td>Paul Richardson</td>\n",
       "      <td>245</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAS</th>\n",
       "      <td>Kelvin Harmon</td>\n",
       "      <td>332</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>WAS</th>\n",
       "      <td>Trey Quinn</td>\n",
       "      <td>198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TEN</th>\n",
       "      <td>A.J. Brown</td>\n",
       "      <td>927</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ATL</th>\n",
       "      <td>Calvin Ridley</td>\n",
       "      <td>866</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ARI</th>\n",
       "      <td>Larry Fitzgerald</td>\n",
       "      <td>759</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ARI</th>\n",
       "      <td>Damiere Byrd</td>\n",
       "      <td>285</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ARI</th>\n",
       "      <td>Pharoh Cooper</td>\n",
       "      <td>219</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ARI</th>\n",
       "      <td>Christian Kirk</td>\n",
       "      <td>649</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>127 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             full_name  rec_yards\n",
       "team                             \n",
       "WAS     Terry McLaurin        919\n",
       "WAS    Paul Richardson        245\n",
       "WAS      Kelvin Harmon        332\n",
       "WAS         Trey Quinn        198\n",
       "TEN         A.J. Brown        927\n",
       "...                ...        ...\n",
       "ATL      Calvin Ridley        866\n",
       "ARI   Larry Fitzgerald        759\n",
       "ARI       Damiere Byrd        285\n",
       "ARI      Pharoh Cooper        219\n",
       "ARI     Christian Kirk        649\n",
       "\n",
       "[127 rows x 2 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.loc[top_n_by_team.index.levels[1].values].set_index('team','full+name').sort_index(ascending=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Speed tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def method1():\n",
    "\n",
    "    # Make a sorted multi-index\n",
    "    df_sorted = df.set_index(['team','rec_yards']).sort_index(ascending=False)\n",
    "\n",
    "    # Add a dummy column containing all 1s\n",
    "    df_sorted['count'] = 1\n",
    "\n",
    "    # Turn it into a ranking by team\n",
    "    df_sorted['count'] = df_sorted.groupby('team')['count'].cumsum()       \n",
    "\n",
    "    # Only keep the first 4 rows per team\n",
    "    df_sorted = df_sorted.loc[df_sorted['count'] <= 4]\n",
    "\n",
    "def method2():\n",
    "    top4_by_team = df.sort_values(['team','rec_yards'], ascending=False).groupby('team').head(4)\n",
    "    \n",
    "def method3():\n",
    "    df.groupby('team').apply(pd.DataFrame.nlargest, 4, columns=['rec_yards'])\n",
    "    \n",
    "def method4():\n",
    "    top4_by_team = df.groupby('team')['rec_yards'].nlargest(4).index.levels[1].values\n",
    "    top4_by_team = df.loc[top4_by_team].set_index('team','full+name').sort_index(ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.34 ms ± 757 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit method1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.02 ms ± 216 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit method2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "62 ms ± 816 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit method3()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21.8 ms ± 754 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit method4()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (torch)",
   "language": "python",
   "name": "torch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}