feat(pandas): add problem 596

Author: IndexSeek

README.md

@@ -39,7 +39,7 @@ Fiddling around with DataFusion, pandas, and PyArrow.
 |         2356 | [Number of Unique Subjects Taught by Each Teacher](https://leetcode.com/problems/number-of-unique-subjects-taught-by-each-teacher) | Easy         | ❌           | ✅       | ✅        |
 |         1141 | [User Activity for the Past 30 Days I](https://leetcode.com/problems/user-activity-for-the-past-30-days-i)                         | Easy         | ❌           | ✅       | ✅        |
 |         1070 | [Product Sales Analysis III](https://leetcode.com/problems/product-sales-analysis-iii)                                             | Medium       | ❌           | ✅       | ✅        |
-|          596 | [Classes More Than 5 Students](https://leetcode.com/problems/classes-more-than-5-students)                                         | Easy         | ❌           | ❌       | ✅        |
+|          596 | [Classes More Than 5 Students](https://leetcode.com/problems/classes-more-than-5-students)                                         | Easy         | ❌           | ✅       | ✅        |
 |         1729 | [Find Followers Count](https://leetcode.com/problems/find-followers-count)                                                         | Easy         | ❌           | ❌       | ✅        |
 |          619 | [Biggest Single Number](https://leetcode.com/problems/biggest-single-number)                                                       | Easy         | ❌           | ❌       | ✅        |
 |         1045 | [Customers Who Bought All Products](https://leetcode.com/problems/customers-who-bought-all-products)                               | Medium       | ❌           | ❌       | ✅        |

problems/pandas.py

@@ -236,6 +236,28 @@ def problem_595(world: pd.DataFrame) -> pd.DataFrame:
     return world[columns]
 
 
+def problem_596(courses: pd.DataFrame) -> pd.DataFrame:
+    """Find all the classes that have at least five students.
+
+    Return the result table in any order.
+
+    Parameters
+    ----------
+    courses : pd.DataFrame
+        Table indicates the name of a student and the class in which they are enrolled.
+
+    Returns
+    -------
+    pd.DataFrame
+
+    """
+    return (
+        courses.groupby("class", as_index=False)
+        .aggregate(count=pd.NamedAgg("student", "count"))
+        .query("count >= 5")[["class"]]
+    )
+
+
 def problem_1068(sales: pd.DataFrame, product: pd.DataFrame) -> pd.DataFrame:
     """Report the product_name, year, and price for each sale_id in the Sales table.
 

tests/test_pandas.py

@@ -433,6 +433,72 @@ def test_problem_595(input_data, expected_data):
     )
 
 
+@pytest.mark.parametrize(
+    "input_data, expected_data",
+    [
+        pytest.param(
+            {
+                "student": [
+                    "A",
+                    "B",
+                    "C",
+                    "D",
+                    "E",
+                    "F",
+                    "G",
+                    "H",
+                    "I",
+                    "J",
+                    "K",
+                    "L",
+                ],
+                "class": [
+                    "Math",
+                    "Math",
+                    "Biology",
+                    "Math",
+                    "Computer",
+                    "Math",
+                    "Math",
+                    "Math",
+                    "Computer",
+                    "Computer",
+                    "Computer",
+                    "Computer",
+                ],
+            },
+            {"class": ["Math", "Computer"]},
+            id="happy_path_multiple_classes",
+        ),
+        pytest.param(
+            {
+                "class": ["History", "History", "History", "History", "History"],
+                "student": ["Alice", "Bob", "Charlie", "David", "Eve"],
+            },
+            {"class": ["History"]},
+            id="edge_case_exactly_5_students",
+        ),
+        pytest.param(
+            {
+                "class": ["Art", "Art", "Art", "Art"],
+                "student": ["Alice", "Bob", "Charlie", "David"],
+            },
+            {"class": []},
+            id="edge_case_less_than_5_students",
+        ),
+    ],
+)
+def test_problem_596(input_data, expected_data):
+    table = pd.DataFrame(input_data)
+    expected_table = (
+        pd.DataFrame(expected_data).sort_values("class").reset_index(drop=True)
+    )
+    result = problem_596(table).sort_values("class").reset_index(drop=True)
+    assert_frame_equal(
+        result, expected_table, check_dtype=False, check_index_type=False
+    )
+
+
 @pytest.mark.parametrize(
     "input_data, expected_data",
     [