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:
    # Group by the specified column, sort within each group by the value column in descending order
    # and limit the number of rows to n. Then, reset the index to get a clean RangeIndex.
    result = df.groupby(group_col, group_keys=False).apply(
        lambda x: x.sort_values(value_col, ascending=False).head(n)
    ).reset_index(drop=True)
    
    return result
```

[
  {
    "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    # Group by the specified column, sort within each group by the value column in descending order\n    # and limit the number of rows to n. Then, reset the index to get a clean RangeIndex.\n    result = df.groupby(group_col, group_keys=False).apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    ).reset_index(drop=True)\n    \n    return result\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    # Group by the specified column, sort within each group by the value column in descending order\n    # and limit the number of rows to n. Then, reset the index to get a clean RangeIndex.\n    # Ensure that the result preserves all original columns.\n    result = df.groupby(group_col, as_index=False)[[group_col, value_col]].apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    ).reset_index(drop=True)\n    \n    # Add back any extra columns that were not in the groupby columns\n    extra_cols = df.columns.difference([group_col, value_col])\n    result = pd.merge(result, df.loc[:, extra_cols], on=group_col, how='left')\n    \n    return result\n```",
    "passed": 0,
    "test_error": "^^^^^^^^^^^^\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \n\nself =   extra\n0     p\n1     q\n2     r, key = 'team', axis = 0\n\n    @final\n    def _get_label_or_level_values(self, key: Level, axis: AxisInt = 0) -> ArrayLike:\n        \"\"\"\n        Return a 1-D array of values associated with `key`, a label or level\n        from the given `axis`.\n    \n        Retrieval logic:\n          - (axis=0): Return column values if `key` matches a column label.\n            Otherwise return index level values if `key` matches an index\n            level.\n          - (axis=1): Return row values if `key` matches an index label.\n            Otherwise return column level values if 'key' matches a column\n            level\n    \n        Parameters\n        ----------\n        key : Hashable\n            Label or level name.\n        axis : int, default 0\n            Axis that levels are associated with (0 for index, 1 for columns)\n    \n        Returns\n        -------\n        np.ndarray or ExtensionArray\n    \n        Raises\n        ------\n        KeyError\n            if `key` matches neither a label nor a level\n        ValueError\n            if `key` matches multiple labels\n        \"\"\"\n        axis = self._get_axis_number(axis)\n        first_other_axes = next(\n            (ax for ax in range(self._AXIS_LEN) if ax != axis), None\n        )\n    \n        if self._is_label_reference(key, axis=axis):\n            self._check_label_or_level_ambiguity(key, axis=axis)\n            if first_other_axes is None:\n                raise ValueError(\"axis matched all axes\")\n            values = self.xs(key, axis=first_other_axes)._values\n        elif self._is_level_reference(key, axis=axis):\n            values = self.axes[axis].get_level_values(key)._values\n        else:\n>           raise KeyError(key)\nE           KeyError: 'team'\n\n/home/psc/miniconda3/lib/python3.12/site-packages/pandas/core/generic.py:1776: KeyError\n=========================== short test summary info ============================\nFAILED test_top_n.py::test_basic_top2 - KeyError: 'team'\nFAILED test_top_n.py::test_index_is_rangeindex - KeyError: 'team'\nFAILED test_top_n.py::test_group_smaller_than_n - KeyError: 'team'\nFAILED test_top_n.py::test_n_equals_one - KeyError: 'team'\nFAILED test_top_n.py::test_does_not_mutate_input - KeyError: 'team'\nFAILED test_top_n.py::test_group_order_follows_first_appearance - KeyError: 'g'\nFAILED test_top_n.py::test_preserves_extra_columns - KeyError: 'team'\n7 failed in 0.46s"
  },
  {
    "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    # Group by the specified column, sort within each group by the value column in descending order\n    # and limit the number of rows to n. Then, reset the index to get a clean RangeIndex.\n    # Ensure that the result preserves all original columns.\n    result = df.groupby(group_col, as_index=False)[[group_col, value_col]].apply(\n        lambda x: x.sort_values(value_col, ascending=False).head(n)\n    ).reset_index(drop=True)\n    \n    # Add back any extra columns that were not in the groupby columns\n    extra_cols = df.columns.difference([group_col, value_col])\n    result = pd.merge(result, df.loc[:, extra_cols], on=group_col, how='left')\n    \n    return result\n```",
    "passed": 0,
    "test_error": ""
  }
]