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:
    # Create a helper column to track the order of groups based on their first appearance.
    # factorize() returns codes in the order of first appearance.
    df_sorted = df.assign(_group_order=df[group_col].factorize()[0])
    
    # Sort by group order (ascending) and value column (descending).
    df_sorted = df_sorted.sort_values(
        by=['_group_order', value_col], 
        ascending=[True, False]
    )
    
    # Group by the specified column (sort=False preserves the order from df_sorted)
    # and take the top n rows for each group.
    result = df_sorted.groupby(group_col, sort=False).head(n)
    
    # Remove the helper column and reset the index.
    return result.drop(columns=['_group_order']).reset_index(drop=True)
```

The user wants a Python function `top_n_per_group` that takes a pandas DataFrame and returns the top N rows per group based on a value column.

Requirements:
1. Input: `df` (DataFrame), `group_col` (str), `value_col` (str), `n` (int).
2. Output: DataFrame with top `n` rows per group by `value_col` descending.
3. Preserve all original columns.
4. Order within group: `value_col` descending.
5. Order of groups: First appearance in `df`.
6. Ties: Arbitrary, but count must be `min(n, group_size)`.
7. Index: Reset to RangeIndex.
8. No mutation of input `df`.

Algorithm:
1. Sort the DataFrame by `group_col` and `value_col`.
   - To maintain group order by first appearance, we can't just sort by `group_col` alphabetically.
   - We need to sort by `value_col` descending within each group, but keep the groups in their original order.
   - One way:
     - Create a helper column or use `sort_values` with a key, or use `groupby` with `sort=False`.
     - `df.sort_values(by=[group_col, value_col], ascending=[True, False])` sorts groups alphabetically (or by type), which might violate "Group order follows the order in which each group first appears in df".
     - To preserve group order:
       - We can use `df.groupby(group_col, sort=False)`. This preserves the order of groups as they appear in the index.
       - However, `sort_values` is usually needed to rank rows.
       - If we do `df.sort_values(value_col, ascending=False)`, the groups get mixed up.
       - If we do `df.sort_values([group_col, value_col], ascending=[True, False])`, groups are sorted by `group_col` value, not appearance.

   - Correct approach for group order:
     - We need to sort rows such that:
       1. Groups appear in order of first appearance.
       2. Within groups, rows are sorted by `value_col` descending.
     - We can achieve this by:
       - Using `df.sort_values(value_col, ascending=False)`? No, that mixes groups.
       - Using `df.groupby(group_col, sort=False).apply(lambda g: g.sort_values(value_col, ascending=False))`?
         - This works but might be slow for large DataFrames.
         - Also, `apply` might mess up indices or column types sometimes, though usually fine.
       - Better approach:
         - Sort by `value_col` descending.
         - Then use `groupby(group_col, sort=False).head(n)`.
         - Wait, if we sort by `value_col` descending globally, the groups are interleaved. `groupby(...).head(n)` will pick the first `n` rows of each group *in the current order*.
         - If the DataFrame is sorted by `value_col` descending globally, the first `n` rows of group 'a' might not be the top `n` of 'a' if 'a' is scattered? No, `groupby` collects all rows of 'a'. If the whole DF is sorted by score desc, then all rows of 'a' are sorted by score desc relative to each other?
         - Yes. If `df` is sorted by `value_col` descending, then for any subset of rows (like those in group 'a'), they are also sorted by `value_col` descending.
         - So:
           1. `df_sorted = df.sort_values(value_col, ascending=False)`
           2. `result = df_sorted.groupby(group_col, sort=False).head(n)`
           3. `result = result.reset_index(drop=True)`
         - Does this preserve group order?
           - `groupby(..., sort=False)` iterates groups in the order they appear in the index of `df_sorted`.
           - Since `df_sorted` is a permutation of `df`, the order of first appearance of groups in `df_sorted` might differ from `df`.
           - Example:
             - df: group 'a' (row 0), group 'b' (row 1).
             - sort_values(score desc): maybe 'b' has higher score, so 'b' comes first.
             - groupby order becomes 'b', then 'a'.
             - Requirement: "Group order follows the order in which each group first appears in df".
             - So this simple sort + groupby approach fails the group order requirement.

   - Refined approach:
     - We need to sort within groups, but keep groups in original order.
     - We can use `df.sort_values(value_col, ascending=False)`? No.
     - We can use `df.groupby(group_col, sort=False).apply(lambda x: x.sort_values(value_col, ascending=False))`.
       - `sort=False` in groupby ensures groups are processed in order of appearance in `df`.
       - `apply` returns a DataFrame. The index might be hierarchical (group, original_index).
       - We need to reset index.
       - This seems robust.
     - Alternative:
       - Use `df.sort_values([group_col, value_col], ascending=[True, False])`?
         - This sorts groups by `group_col` value. If `group_col` is categorical or string, it sorts alphabetically. This violates "order in which each group first appears".
         - Unless we map `group_col` to an integer based on first appearance.
         - `group_order = df[group_col].drop_duplicates().tolist()`
         - `df['group_rank'] = df[group_col].map({g: i for i, g in enumerate(group_order)})`
         - `df_sorted = df.sort_values(['group_rank', value_col], ascending=[True, False])`
         - `result = df_sorted.groupby(group_col, sort=False).head(n)`
         - `result = result.drop(columns=['group_rank']).reset_index(drop=True)`
         - This works and is efficient.

   - Let's check the `apply` method again.
     - `df.groupby(group_col, sort=False).apply(lambda g: g.sort_values(value_col, ascending=False))`
     - This creates a new DataFrame.
     - The index will be MultiIndex (group_key, original_index).
     - `reset_index(drop=True)` will flatten it.
     - Is it efficient? For very large DataFrames, `apply` can be slower than vectorized operations, but for this task, it's likely acceptable.
     - However, `sort_values` inside `apply` is essentially doing what we want.
     - One issue: `apply` might infer dtypes differently or have overhead.
     - The `group_rank` approach is purely vectorized sorting.

   - Let's look at the requirements again.
     - "Group order follows the order in which each group first appears in df."
     - "Within each group, rows are ordered by value_col descending."
     - "Ties ... arbitrary".
     - "Clean RangeIndex".

   - Let's try th
…[22189 chars elided]…
ecessary? `df.copy()` then `df_copy['_group_order'] = ...` is also fine. `assign` is just cleaner.
   - I'll stick with `assign`.

   - One check: `df[group_col].factorize()` might raise if `group_col` is not in `df`. But that's expected.
   - What if `df` is empty? `factorize` returns empty. `sort_values` returns empty. `groupby.head` returns empty. `drop` returns empty. `reset_index` returns empty. Correct.

   - Final check on "Do not mutate the input df".
     - `df.assign` returns a new DataFrame.
     - `df_sorted` is a new DataFrame.
     - `df_sorted.sort_values` returns a new DataFrame.
     - `df_sorted.groupby(...).head(n)` returns a new DataFrame.
     - `result.drop(...)` returns a new DataFrame.
     - `result.reset_index(...)` returns a new DataFrame.
     - `df` is never mutated.

   - Looks perfect.

```python
import pandas as pd

def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:
    # Create a helper column to track the order of groups based on their first appearance.
    # factorize() returns codes in the order of first appearance.
    df_sorted = df.assign(_group_order=df[group_col].factorize()[0])
    
    # Sort by group order (ascending) and value column (descending).
    df_sorted = df_sorted.sort_values(
        by=['_group_order', value_col], 
        ascending=[True, False]
    )
    
    # Group by the specified column (sort=False preserves the order from df_sorted)
    # and take the top n rows for each group.
    result = df_sorted.groupby(group_col, sort=False).head(n)
    
    # Remove the helper column and reset the index.
    return result.drop(columns=['_group_order']).reset_index(drop=True)
```
   - I'll remove comments for brevity if needed, but they are helpful. The prompt doesn't restrict comments.
   - I'll provide the code block.

   - One tiny thing: `df[group_col].factorize()` vs `pd.factorize(df[group_col])`. Both are fine. `Series.factorize` is slightly more direct.
   - I'll use `df[group_col].factorize()[0]`.

   - Ready.