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:
    """
    Returns the top n rows per group defined by group_col, ranked by value_col in descending order.
    
    Args:
        df: Input DataFrame (not mutated).
        group_col: Column name to group by.
        value_col: Column name to rank by.
        n: Number of top rows to return per group.
        
    Returns:
        A new DataFrame with the top n rows per group, ordered by group and value_col descending.
    """
    # Sort by group first to ensure group order follows first appearance
    sorted_df = df.sort_values(group_col)
    
    # Group by group_col, sort by value_col descending, and take the top n
    result = sorted_df.groupby(group_col, sort=False).apply(
        lambda x: x.sort_values(value_col, ascending=False).head(n)
    )
    
    # Reset index to ensure clean RangeIndex and drop the old group index
    return result.reset_index(drop=True)
```

[
  {
    "total": 7,
    "answer": "```python\nimport pandas as pd\n\ndef top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:\n    \"\"\"\n    Returns the top n rows per group defined by group_col, ranked by value_col in descending order.\n    \n    Args:\n        df: Input DataFrame (not mutated).\n        group_col: Column name to group by.\n        value_col: Column name to rank by.\n        n: Number of top rows to return per group.\n        \n    Returns:\n        A new DataFrame with the top n rows per group, ordered by group and value_col descending.\n    \"\"\"\n    # Sort by group first to ensure group order follows first appearance\n    sorted_df = df.sort_values(group_col)\n    \n    # Group by group_col, sort by value_col descending, and take the top n\n    result = sorted_df.groupby(group_col, sort=False).apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    )\n    \n    # Reset index to ensure clean RangeIndex and drop the old group index\n    return result.reset_index(drop=True)\n```",
    "passed": 3,
    "test_error": "          Integer location, slice or boolean mask.\n    \n        See Also\n        --------\n        Index.get_slice_bound : Calculate slice bound that corresponds to\n            given label.\n        Index.get_indexer : Computes indexer and mask for new index given\n            the current index.\n        Index.get_non_unique : Returns indexer and masks for new index given\n            the current index.\n        Index.get_indexer_for : Returns an indexer even when non-unique.\n    \n        Examples\n        --------\n        >>> unique_index = pd.Index(list(\"abc\"))\n        >>> unique_index.get_loc(\"b\")\n        1\n    \n        >>> monotonic_index = pd.Index(list(\"abbc\"))\n        >>> monotonic_index.get_loc(\"b\")\n        slice(1, 3, None)\n    \n        >>> non_monotonic_index = pd.Index(list(\"abcb\"))\n        >>> non_monotonic_index.get_loc(\"b\")\n        array([False,  True, False,  True])\n        \"\"\"\n        casted_key = self._maybe_cast_indexer(key)\n        try:\n            return self._engine.get_loc(casted_key)\n        except KeyError as err:\n            if isinstance(casted_key, slice) or (\n                isinstance(casted_key, abc.Iterable)\n                and any(isinstance(x, slice) for x in casted_key)\n            ):\n                raise InvalidIndexError(key) from err\n>           raise KeyError(key) from err\nE           KeyError: 'g'\n\n/home/psc/miniconda3/lib/python3.12/site-packages/pandas/core/indexes/base.py:3648: KeyError\n_________________________ test_preserves_extra_columns _________________________\n\n    def test_preserves_extra_columns():\n        df = pd.DataFrame({\n            \"team\": [\"a\", \"a\", \"b\"],\n            \"score\": [1, 2, 3],\n            \"extra\": [\"p\", \"q\", \"r\"],\n        })\n        out = top_n_per_group(df, \"team\", \"score\", 1)\n>       assert list(out.columns) == [\"team\", \"score\", \"extra\"]\nE       AssertionError: assert ['score', 'extra'] == ['team', 'score', 'extra']\nE         \nE         At index 0 diff: 'score' != 'team'\nE         Right contains one more item: 'extra'\nE         Use -v to get more diff\n\ntest_top_n.py:72: AssertionError\n=========================== short test summary info ============================\nFAILED test_top_n.py::test_basic_top2 - AssertionError: assert ['name', 'scor...\nFAILED test_top_n.py::test_group_smaller_than_n - KeyError: 'team'\nFAILED test_top_n.py::test_group_order_follows_first_appearance - KeyError: 'g'\nFAILED test_top_n.py::test_preserves_extra_columns - AssertionError: assert [...\n4 failed, 3 passed in 0.33s"
  },
  {
    "total": 7,
    "answer": "```python\nimport pandas as pd\n\ndef top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:\n    \"\"\"\n    Returns the top n rows per group defined by group_col, ranked by value_col in descending order.\n    \n    Args:\n        df: Input DataFrame (not mutated).\n        group_col: Column name to group by.\n        value_col: Column name to rank by.\n        n: Number of top rows to return per group.\n        \n    Returns:\n        A new DataFrame with the top n rows per group, ordered by group and value_col descending.\n    \"\"\"\n    # Sort by group first to ensure group order follows first appearance\n    sorted_df = df.sort_values(group_col)\n    \n    # Group by group_col, sort by value_col descending, and take the top n\n    result = sorted_df.groupby(group_col, sort=False).apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    )\n    \n    # Reset index to ensure clean RangeIndex and drop the old group index\n    return result.reset_index(drop=True)\n```",
    "passed": 3,
    "test_error": "          Integer location, slice or boolean mask.\n    \n        See Also\n        --------\n        Index.get_slice_bound : Calculate slice bound that corresponds to\n            given label.\n        Index.get_indexer : Computes indexer and mask for new index given\n            the current index.\n        Index.get_non_unique : Returns indexer and masks for new index given\n            the current index.\n        Index.get_indexer_for : Returns an indexer even when non-unique.\n    \n        Examples\n        --------\n        >>> unique_index = pd.Index(list(\"abc\"))\n        >>> unique_index.get_loc(\"b\")\n        1\n    \n        >>> monotonic_index = pd.Index(list(\"abbc\"))\n        >>> monotonic_index.get_loc(\"b\")\n        slice(1, 3, None)\n    \n        >>> non_monotonic_index = pd.Index(list(\"abcb\"))\n        >>> non_monotonic_index.get_loc(\"b\")\n        array([False,  True, False,  True])\n        \"\"\"\n        casted_key = self._maybe_cast_indexer(key)\n        try:\n            return self._engine.get_loc(casted_key)\n        except KeyError as err:\n            if isinstance(casted_key, slice) or (\n                isinstance(casted_key, abc.Iterable)\n                and any(isinstance(x, slice) for x in casted_key)\n            ):\n                raise InvalidIndexError(key) from err\n>           raise KeyError(key) from err\nE           KeyError: 'g'\n\n/home/psc/miniconda3/lib/python3.12/site-packages/pandas/core/indexes/base.py:3648: KeyError\n_________________________ test_preserves_extra_columns _________________________\n\n    def test_preserves_extra_columns():\n        df = pd.DataFrame({\n            \"team\": [\"a\", \"a\", \"b\"],\n            \"score\": [1, 2, 3],\n            \"extra\": [\"p\", \"q\", \"r\"],\n        })\n        out = top_n_per_group(df, \"team\", \"score\", 1)\n>       assert list(out.columns) == [\"team\", \"score\", \"extra\"]\nE       AssertionError: assert ['score', 'extra'] == ['team', 'score', 'extra']\nE         \nE         At index 0 diff: 'score' != 'team'\nE         Right contains one more item: 'extra'\nE         Use -v to get more diff\n\ntest_top_n.py:72: AssertionError\n=========================== short test summary info ============================\nFAILED test_top_n.py::test_basic_top2 - AssertionError: assert ['name', 'scor...\nFAILED test_top_n.py::test_group_smaller_than_n - KeyError: 'team'\nFAILED test_top_n.py::test_group_order_follows_first_appearance - KeyError: 'g'\nFAILED test_top_n.py::test_preserves_extra_columns - AssertionError: assert [...\n4 failed, 3 passed in 0.28s"
  },
  {
    "total": 7,
    "answer": "```python\nimport pandas as pd\n\ndef top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:\n    \"\"\"\n    Returns the top n rows per group defined by group_col, ranked by value_col in descending order.\n    \n    Args:\n        df: Input DataFrame (not mutated).\n        group_col: Column name to group by.\n        value_col: Column name to rank by.\n        n: Number of top rows to return per group.\n        \n    Returns:\n        A new DataFrame with the top n rows per group, ordered by group and value_col descending.\n    \"\"\"\n    # Sort by group first to ensure group order follows first appearance\n    sorted_df = df.sort_values(group_col)\n    \n    # Group by group_col, sort by value_col descending, and take the top n\n    result = sorted_df.groupby(group_col, sort=False).apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    )\n    \n    # Reset index to ensure clean RangeIndex and drop the old group index\n    return result.reset_index(drop=True)\n```",
    "passed": 3,
    "test_error": ""
  }
]