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.

Looking at the requirements, I need to implement a function that:

1. Groups the DataFrame by `group_col`
2. Within each group, sorts by `value_col` in descending order
3. Takes the top `n` rows from each group (or all rows if group has fewer than `n`)
4. Preserves the original column order and all columns
5. Orders groups by their first appearance in the original DataFrame
6. Returns a clean RangeIndex

The key challenge is maintaining the group order as they first appear in the original DataFrame. I can achieve this by using `groupby` with `sort=False` (which preserves original order) or by tracking the order explicitly.

Here's my implementation:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Sort by group_col (to maintain original group order) and then by value_col descending
    # First, we need to preserve the original order of groups
    # We can do this by adding a temporary column with original index, then sort appropriately
    
    # Add a temporary column to track original row order
    df_copy['_original_order'] = range(len(df_copy))
    
    # Sort by group_col (preserving first appearance order), then by value_col descending
    # For groups, we want to maintain the order they first appear in the original df
    # We can achieve this by sorting by group_col with sort=False, but we need to handle
    # the value_col sorting within groups
    
    # Alternative approach: use groupby with sort=False to preserve group order
    # First, sort within each group by value_col descending
    sorted_df = df_copy.sort_values([group_col, value_col], ascending=[True, False])
    
    # Group by group_col and take top n from each group
    result = sorted_df.groupby(group_col, sort=False).head(n)
    
    # Remove the temporary column
    result = result.drop(columns=['_original_order'])
    
    # Reset index to get clean RangeIndex
    result = result.reset_index(drop=True)
    
    return result
```

Wait, let me reconsider. The `sort_values([group_col, value_col], ascending=[True, False])` approach might not preserve the original group order correctly because `sort_values` will sort the groups alphabetically/numerically by default.

Let me use a more reliable approach:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Add a temporary column to track original row order
    df_copy['_original_order'] = range(len(df_copy))
    
    # Sort by group_col preserving first appearance order, then by value_col descending
    # To preserve group order, we need to sort by the first occurrence of each group
    # We can use groupby with sort=False which preserves the order groups appear in the data
    
    # First, sort within each group by value_col descending
    # Use groupby with sort=False to maintain original group order
    result_list = []
    
    for name, group in df_copy.groupby(group_col, sort=False):
        # Sort this group by value_col descending
        sorted_group = group.sort_values(value_col, ascending=False)
        # Take top n rows
        top_n = sorted_group.head(n)
        result_list.append(top_n)
    
    # Concatenate all groups
    result = pd.concat(result_list, ignore_index=True)
    
    # Remove the temporary column
    result = result.drop(columns=['_original_order'])
    
    return result
```

Actually, there's a simpler approach using pandas' built-in functionality. The `groupby().head(n)` method works well, but we need to ensure groups are processed in their first-appearance order.

Here's the cleanest solution:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Sort by value_col descending within each group
    # First sort by group_col (to maintain first appearance order) and then by value_col descending
    # But we need to be careful about group order
    
    # Use groupby with sort=False to preserve original group order
    # Sort the dataframe by group_col and value_col, but ensure groups maintain their first-appearance order
    sorted_df = df_copy.sort_values([group_col, value_col], ascending=[True, False])
    
    # Group by group_col without sorting (preserves first appearance order)
    result = sorted_df.groupby(group_col, sort=False).head(n)
    
    # Reset index to get clean RangeIndex
    result = result.reset_index(drop=True)
    
    return result
```

Wait, I think the issue is that `sort_values([group_col, value_col], ascending=[True, False])` will sort groups alphabetically, not by first appearance. Let me fix this properly.

The correct approach is to use `groupby` with `sort=False` which preserves the order groups appear in the original DataFrame:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Sort by value_col descending first (to get highest values at top within each group)
    sorted_df = df_copy.sort_values(value_col, ascending=False)
    
    # Group by group_col without sorting (preserves first appearance order in original df)
    result = sorted_df.groupby(group_col, sort=False).head(n)
    
    # Reset index to get clean RangeIndex
    result = result.reset_index(drop=True)
    
    return result
```

Let me test this with the example:
- Original df: team a appears first (rows 0,1,2), then team b (rows 3,4)
- After sorting by score descending: y(30), q(15), z(20), p(5), x(10) - wait, this would put q before z
- Then groupby with sort=False would process groups in order of first appearance in the sorted df, not original df

I need to preserve the original group order from the input DataFrame. Let me use a different approach:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Add a column to track the first appearance order of each group
    # Get unique groups in order of first appearance
    group_order = df_copy[group_col].drop_duplicates().reset_index(drop=True)
    group_order.name = '_group_order'
    
    # Merge to add group order
    df_with_order = df_copy.merge(group_order.reset_index(), left_on=group_col, right_on=group_col, how='left')
    
    # Sort by group order, then by value_col descending
    sorted_df = df_with_order.sort_values(['_group_order', value_col], ascending=[True, False])
    
    # Group by group_col and take top n
    result = sorted_df.groupby(group_col, sort=False).head(n)
    
    # Drop the helper column
    result = result.drop(columns=['_group_order'])
    
    # Reset index
    result = result.reset_index(drop=True)
    
    return result
```

Actually, there's a much simpler way. I can use `groupby` with `sort=False` and apply a custom function:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Sort by value_col descending first
    sorted_df = df_copy.sort_values(value_col, ascending=False)
    
    # Group by group_col without sorting (preserves first appearance order in original df)
    result = sorted_df.groupby(group_col, sort=False).head(n)
    
    # Reset index to get clean RangeIndex
    result = result.reset_index(drop=True)
    
    return result
```

Wait, I need to be more careful. The `sort=False` parameter in `groupby` preserves the order groups appear in the data being grouped. So if I sort by value_col first, the group order will be based on the first appearance in that sorted data, not the original data.

Let me implement it correctly:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Add a column to track original row order
    df_copy['_original_order'] = range(len(df_copy))
    
    # Sort by group_col (preserving first appearance order), then by value_col descending
    # To preserve group order, we need to sort by the first occurrence of each group
    # We can use groupby with sort=False after sorting within groups
    
    # First, sort by value_col descending to get highest values at top
    sorted_by_value = df_copy.sort_values([value_col], ascending=False)
    
    # Now group by group_col without sorting to preserve first appearance order in original df
    # But we need to ensure that within each group, rows are ordered by value_col descending
    
    # The cleanest approach: use groupby with sort=False on the original df after sorting by value_col
    # Actually, let's use the approach from pandas documentation
    
    result = (df_copy
             .sort_values([value_col], ascending=False)
             .groupby(group_col, sort=False)
             .head(n)
             .sort_values([group_col, '_original_order'])  # Sort by group and original order
             .drop(columns=['_original_order'])
             .reset_index(drop=True))
    
    return result
```

I'm overcomplicating this. Let me look at the example again and implement correctly:

```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 copy to avoid mutating the input
    df_copy = df.copy()
    
    # Add a column to track original row order for final sorting
    df_copy['_original_order'] = range(len(df_copy))
    
    # Sort by value_col descending first (to get highest values at top)
    sorted_df = df_copy.sort_values(value_col, ascending=False)
    
    # Group by group_col without sorting (preserves first appearance order in original df)
    grouped = sorted_df.groupby(group_col, sort=False)
    
    # Take top n from each group
    result = grouped.head(n)
    
    # Sort the final result by group_col (preserving first appearance order) and value_col descending
    # We need to maintain the order: groups in first-appearance order, within groups sorted by value_col desc
    
    # Get unique groups in order of first appearance
    group_order = df_copy[group_col].drop_duplicates().reset_index(drop=True)
    
    # Create a mapping from group to its order
    group_order_map = {group: i for i, group in enumerate(group_order)}
    
    # Add group order to result
    result['_group_order'] = result[group_col].map(group_order_map)
    
    # Sort by group order and value_col descending
    result = result.sort_values(['_group_order', value_col], ascending=[True, False])
    
    # Drop helper columns and reset index
    result = result.drop(columns=['_group_order', '_original_order']).reset_index(drop=True)
    
    return result
```