py-07-pandas-top-n
0.857
Challenge · difficulty 4/5
# 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
```
tests/test_top_n.py
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"
Proposed solution
```python
import pandas as pd
def top_n_per_group(df: pd.DataFrame, group_col: str, value_col: str, n: int) -> pd.DataFrame:
return (
df.sort_values(value_col, ascending=False)
.groupby(group_col, sort=False)
.head(n)
.reset_index(drop=True)
)
```Test output (stdout)
.....F. [100%]
=================================== FAILURES ===================================
__________________ test_group_order_follows_first_appearance ___________________
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"]
E AssertionError: assert ['a', 'z'] == ['z', 'a']
E
E At index 0 diff: 'a' != 'z'
E Use -v to get more diff
test_top_n.py:61: AssertionError
=========================== short test summary info ============================
FAILED test_top_n.py::test_group_order_follows_first_appearance - AssertionEr...
1 failed, 6 passed in 0.17s