py-07-pandas-top-n

# Top-N rows per group (pandas)

Implement **`solution.py`** with:

```python
import pandas as pd

def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:
    ...
```

Given a DataFrame `df`, return a new DataFrame containing, for each group defined
by `group_col`, the **top `n` rows ranked by `value_col` in descending order**.

Requirements:

- All original columns must be preserved (do not drop, rename, or reorder columns).
- Within each group, rows are ordered by `value_col` **descending**. Groups that
  have fewer than `n` rows contribute all of their rows.
- The result is ordered by group, and within each group by `value_col` descending.
  Group order follows the order in which each group first appears in `df`.
- Ties in `value_col` may be broken arbitrarily, but the number of rows returned
  per group must be exactly `min(n, group_size)`.
- The returned DataFrame must use a clean `RangeIndex` (`0..len-1`) — call
  `reset_index(drop=True)` on the result.
- Do not mutate the input `df`.

Use pandas (e.g. `sort_values` + `groupby(...).head(n)`).

Example:

```python
df = pd.DataFrame({
    "team": ["a", "a", "a", "b", "b"],
    "name": ["x", "y", "z", "p", "q"],
    "score": [10, 30, 20, 5, 15],
})
top_n_per_group(df, "team", "score", 2)
#   team name  score
# 0    a    y     30
# 1    a    z     20
# 2    b    q     15
# 3    b    p      5
```

import pandas as pd
from solution import top_n_per_group


def base_df():
    return pd.DataFrame({
        "team": ["a", "a", "a", "b", "b"],
        "name": ["x", "y", "z", "p", "q"],
        "score": [10, 30, 20, 5, 15],
    })


def test_basic_top2():
    df = base_df()
    out = top_n_per_group(df, "team", "score", 2)
    assert list(out.columns) == ["team", "name", "score"]
    assert len(out) == 4
    a = out[out["team"] == "a"]
    assert list(a["score"]) == [30, 20]
    assert list(a["name"]) == ["y", "z"]
    b = out[out["team"] == "b"]
    assert list(b["score"]) == [15, 5]


def test_index_is_rangeindex():
    df = base_df()
    out = top_n_per_group(df, "team", "score", 2)
    assert list(out.index) == [0, 1, 2, 3]


def test_group_smaller_than_n():
    df = base_df()
    out = top_n_per_group(df, "team", "score", 10)
    # all rows kept (3 in a, 2 in b)
    assert len(out) == 5
    assert (out["team"] == "a").sum() == 3
    assert (out["team"] == "b").sum() == 2


def test_n_equals_one():
    df = base_df()
    out = top_n_per_group(df, "team", "score", 1)
    assert len(out) == 2
    assert set(out["name"]) == {"y", "q"}


def test_does_not_mutate_input():
    df = base_df()
    snapshot = df.copy()
    _ = top_n_per_group(df, "team", "score", 2)
    pd.testing.assert_frame_equal(df, snapshot)


def test_group_order_follows_first_appearance():
    df = pd.DataFrame({
        "g": ["z", "a", "z", "a"],
        "v": [1, 2, 3, 4],
    })
    out = top_n_per_group(df, "g", "v", 1)
    # "z" appears first, so it should come first
    assert list(out["g"]) == ["z", "a"]
    assert list(out["v"]) == [3, 4]


def test_preserves_extra_columns():
    df = pd.DataFrame({
        "team": ["a", "a", "b"],
        "score": [1, 2, 3],
        "extra": ["p", "q", "r"],
    })
    out = top_n_per_group(df, "team", "score", 1)
    assert list(out.columns) == ["team", "score", "extra"]
    a = out[out["team"] == "a"].iloc[0]
    assert a["extra"] == "q"

You are an expert programmer. Solve the task exactly as specified. Output your solution as fenced code blocks using the required file name(s) and the exact function/type signatures requested. Prefer correctness; do not include prose outside code unless asked.

```python
import pandas as pd

def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:
    # Assign a temporary column to preserve the original group appearance order
    df_temp = df.assign(_group_order=df.groupby(group_col, sort=False).ngroup())
    
    # Sort by group order first, then by value column descending
    df_sorted = df_temp.sort_values(['_group_order', value_col], ascending=[True, False])
    
    # Take top n per group, drop helper column, and reset index
    return df_sorted.groupby(group_col, sort=False).head(n).drop(columns=['_group_order']).reset_index(drop=True)
```

The user wants a function `top_n_per_group` that takes a pandas DataFrame, a group column name, a value column name, and an integer `n`.
It should return the top `n` rows per group based on the value column in descending order.
Requirements:
- Preserve all original columns.
- Order within group: `value_col` descending.
- Order of groups: order of first appearance in `df`.
- Ties: arbitrary, but exactly `min(n, group_size)` rows per group.
- Reset index to `RangeIndex`.
- Do not mutate input `df`.
- Use pandas methods like `sort_values` + `groupby(...).head(n)`.

Let's break down the steps:
1. Sort the DataFrame by `group_col` and `value_col` descending. But wait, we need to preserve the original group order (order of first appearance).
   If we just sort by `group_col`, pandas will sort groups alphabetically or numerically, which might not match the first appearance order.
   To preserve first appearance order, we can assign a rank to each group based on first appearance, or use `sort_values` with a custom key, or simply use `groupby` with `sort=False`.
   Actually, `df.sort_values(['group_col', 'value_col'], ascending=[False, False])` might mess up group order.
   Better approach:
   - Create a temporary column for group order, or use `df.groupby(group_col, sort=False).apply(...)` but `apply` can be slow and might reorder.
   - Actually, `df.sort_values(value_col, ascending=False)` then `groupby(group_col, sort=False).head(n)` works well!
   Let's verify: `sort_values` on `value_col` descending will sort the whole DataFrame. Then `groupby(group_col, sort=False).head(n)` will take the first `n` rows of each group. Since the DataFrame is sorted by `value_col` descending globally, within each group, the rows will be in descending order of `value_col`. And `sort=False` in `groupby` preserves the order of groups as they appear in the DataFrame. Wait, does `groupby(..., sort=False).head(n)` preserve the original group order? Yes, `sort=False` means groups are yielded in the order they appear in the data.
   But wait, if we sort the whole DataFrame by `value_col` descending first, the group order in the sorted DataFrame might change. `groupby` on the sorted DataFrame with `sort=False` will yield groups in the order they appear in the *sorted* DataFrame, not the original.
   So we need to preserve the original group order.

   How to preserve original group order?
   We can add a temporary column `group_order` that assigns an integer to each group based on first appearance.
   Or, we can use `df.sort_values([group_col, value_col], ascending=[False, False])` but that sorts groups.
   Actually, pandas `groupby` with `sort=False` preserves the order of groups as they appear in the *input* to `groupby`. So if we don't sort the DataFrame globally, but sort within groups, that's better.
   We can use `df.sort_values(value_col, ascending=False)`? No.
   Let's think: `df.groupby(group_col, sort=False).apply(lambda g: g.sort_values(value_col, ascending=False)).head(n)`? `apply` might be slow and changes index.
   Better: `df.sort_values(value_col, ascending=False)` then `groupby`? No.
   Actually, we can just do:
   `df.sort_values([group_col, value_col], ascending=[False, False])` -> sorts groups. Not what we want.
   We want groups in order of first appearance.
   We can create a mapping of group values to their first appearance index:
   `group_order = df.groupby(group_col, sort=False).cumcount()`? No.
   `group_order = df[group_col].map(df.groupby(group_col, sort=False).first().index)`? No.
   Simpler: `df['group_rank'] = df.groupby(group_col, sort=False).ngroup()`
   Then sort by `['group_rank', value_col]` ascending `[True, False]`.
   Then drop `group_rank`, groupby `group_col` with `sort=False`, take `head(n)`, reset index.
   Let's test this mentally:
   `df['group_rank'] = df.groupby(group_col, sort=False).ngroup()`
   `sorted_df = df.sort_values(['group_rank', value_col], ascending=[True, False])`
   `result = sorted_df.groupby(group_col, sort=False).head(n)`
   `result = result.drop(columns=['group_rank']).reset_index(drop=True)`
   This preserves group order (by `group_rank`), sorts within group by `value_col` desc, takes top `n`, resets index.
   Does `ngroup()` preserve order? Yes, with `sort=False`, `ngroup()` assigns 0 to the first group encountered, 1 to the second, etc.
   Wait, `df.groupby(group_col, sort=False).ngroup()` works in pandas.
   Let's verify with example:
   df:
   team: a, a, a, b, b
   ngroup() -> 0, 0, 0, 1, 1
   sort_values(['group_rank', 'score'], ascending=[True, False]) ->
   a, y, 30 (rank 0)
   a, z, 20 (rank 0)
   a, x, 10 (rank 0)
   b, q, 15 (rank 1)
   b, p, 5  (rank 1)
   groupby('team', sort=False).head(2) ->
   a, y, 30
   a, z, 20
   b, q, 15
   b, p, 5
   Matches example output exactly.

   Is there a simpler way?
   `df.sort_values(value_col, ascending=False).groupby(group_col, sort=False).head(n)`?
   If we sort globally by `value_col` desc, the order of groups in the sorted df will be determined by which group has the highest values first. `groupby(..., sort=False)` on that will yield groups in that new order, not original. So that's wrong.

   Another way: `df.groupby(group_col, sort=False).apply(lambda x: x.sort_values(value_col, ascending=False).head(n)).reset_index(drop=True)`
   This works but `apply` can be slow and might add an extra index level. We'd need to handle that. Also `apply` is generally discouraged for simple operations.

   The `ngroup()` approach is clean and efficient.
   Let's check pandas version compatibility. `ngroup()` has been around for a long time.
   Alternatively, we can use `df.sort_values(value_col, ascending=False)` and then use `groupby` with `sort=False`? No, as discussed.
   What about `df.sort_values([group_col, value_col], ascending=[False, False])`? It sorts groups alphabetically. The problem says "Group order follows the order in which each group first appears in df." So we must preserve that.

   Let's stick with `ngroup()`.
   Code:
   ```python
   import pandas as pd

   def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) 
…[16787 chars elided]…
(drop=True)
   ```
   This looks solid.
   I will present it. 
   One minor point: `df.groupby(...).ngroup()` might raise a warning if `group_col` contains NaNs, but typically in these problems we assume valid keys. If needed, `pd.factorize` handles NaNs by assigning -1, which would sort them first. `ngroup` also handles it. I'll stick with `ngroup`.
   All good.
   Output matches requirements.
   Proceeds. 
   Self-Correction/Refinement during thought: 
   Actually, `df.sort_values(value_col, ascending=False).groupby(group_col, sort=False).head(n)` is so commonly expected that maybe the group order requirement is just to say "don't sort groups". But to be strictly correct, I'll keep the `_group_order` approach. It's safer.
   I'll output the solution. 
   One thing: `df.assign` might create a copy anyway. It's fine.
   Final code matches the block above.
   Ready. 
   Output generation. 
   [Done] 
   `solution.py` content:
   ```python
   import pandas as pd

   def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:
       df_temp = df.assign(_group_order=df.groupby(group_col, sort=False).ngroup())
       df_sorted = df_temp.sort_values(['_group_order', value_col], ascending=[True, False])
       return df_sorted.groupby(group_col, sort=False).head(n).drop(columns=['_group_order']).reset_index(drop=True)
   ```
   Checks:
   - Preserves columns: yes, drops only temp col.
   - Order within group: value_col desc: yes.
   - Group order: first appearance: yes, via ngroup + sort_values.
   - Ties: arbitrary, min(n, size): yes, head(n) handles it.
   - RangeIndex: yes, reset_index(drop=True).
   - No mutation: yes, assign creates new df.
   - Uses sort_values + groupby.head: yes.
   All constraints satisfied.
   Output matches.✅
   I will produce the response. 
   One minor consideration: `df.groupby(group_col, sort=False).ngroup()` might be slightly slower than `pd.factorize`, but it's fine. I'll stick with it.
   Final answer generation. 
   [Output Generation] -> *Proceeds*