feat(pandas): add problem 1907

Author: IndexSeek

README.md

@@ -52,7 +52,7 @@ Fiddling around with DataFusion, pandas, and PyArrow.
 |          180 | [Consecutive Numbers](https://leetcode.com/problems/consecutive-numbers)                                                                     | Medium       | ✅           | ✅       | ✅        |
 |         1164 | [Product Price at a Given Date](https://leetcode.com/problems/product-price-at-a-given-date)                                                 | Medium       | ❌           | ❌       | ✅        |
 |         1204 | [Last Person to Fit in the Bus](https://leetcode.com/problems/last-person-to-fit-in-the-bus)                                                 | Medium       | ❌           | ❌       | ✅        |
-|         1907 | [Count Salary Categories](https://leetcode.com/problems/count-salary-categories)                                                             | Medium       | ❌           | ❌       | ✅        |
+|         1907 | [Count Salary Categories](https://leetcode.com/problems/count-salary-categories)                                                             | Medium       | ❌           | ✅       | ✅        |
 ## Subqueries
 |   problem_id | title                                                                                                                     | difficulty   | DataFusion   | pandas   | PyArrow   |
 |-------------:|:--------------------------------------------------------------------------------------------------------------------------|:-------------|:-------------|:---------|:----------|

problems/pandas.py

@@ -1258,6 +1258,53 @@ def problem_1789(employee: pd.DataFrame) -> pd.DataFrame:
     ]
 
 
+def problem_1907(accounts: pd.DataFrame) -> pd.DataFrame:
+    """Calculate the number of bank accounts for each salary category.
+
+    The salary categories are:
+
+    - "Low Salary": All the salaries strictly less than $20000.
+    - "Average Salary": All the salaries in the inclusive range [$20000, $50000].
+    - "High Salary": All the salaries strictly greater than $50000.
+
+    The result table must contain all three categories. If there are no accounts in a category, return 0.
+
+    Return the result table in any order.
+
+    Parameters
+    ----------
+    accounts : pd.DataFrame
+        A table containing the account data.
+
+    Returns
+    -------
+    pd.DataFrame
+
+    """
+    accounts["category"] = accounts["income"].case_when(
+        [
+            (accounts["income"] < 20_000, "Low Salary"),
+            (
+                (accounts["income"] >= 20_000) & (accounts["income"] <= 50_000),
+                "Average Salary",
+            ),
+            (accounts["income"] > 50_000, "High Salary"),
+        ]
+    )
+    return (
+        pd.DataFrame(
+            ["Low Salary", "Average Salary", "High Salary"], columns=["category"]
+        )
+        .merge(
+            accounts.groupby("category", as_index=False).aggregate(
+                accounts_count=pd.NamedAgg("account_id", "count")
+            ),
+            how="left",
+        )
+        .fillna(0)
+    )
+
+
 def problem_1934(signups: pd.DataFrame, confirmations: pd.DataFrame) -> pd.DataFrame:
     """Find the confirmation rate of each user.
 

tests/test_pandas.py

@@ -2293,6 +2293,36 @@ def test_problem_1789(input_data, expected_data):
     )
 
 
+@pytest.mark.parametrize(
+    "input_data, expected_data",
+    [
+        pytest.param(
+            {"account_id": [1, 2, 3], "income": [10000, 25000, 100000]},
+            {
+                "category": ["Low Salary", "Average Salary", "High Salary"],
+                "accounts_count": [1, 1, 1],
+            },
+            id="all_three_salary_categories",
+        ),
+        pytest.param(
+            {"account_id": [1, 2], "income": [10000, 25000]},
+            {
+                "category": ["Low Salary", "Average Salary", "High Salary"],
+                "accounts_count": [1, 1, 0],
+            },
+            id="missing_one_salary_category",
+        ),
+    ],
+)
+def test_problem_1907(input_data, expected_data):
+    table = pd.DataFrame(input_data)
+    expected_table = pd.DataFrame(expected_data)
+    result = problem_1907(table).reset_index(drop=True)
+    assert_frame_equal(
+        result, expected_table, check_dtype=False, check_index_type=False
+    )
+
+
 @pytest.mark.parametrize(
     "input_data_1, input_data_2, expected_data",
     [