feat(pandas): add problem 550

Author: IndexSeek

README.md

@@ -32,7 +32,7 @@ Fiddling around with DataFusion, pandas, and PyArrow.
 |         1211 | [Queries Quality and Percentage](https://leetcode.com/problems/queries-quality-and-percentage)                 | Easy         | ❌           | ✅       | ✅        |
 |         1193 | [Monthly Transactions I](https://leetcode.com/problems/monthly-transactions-i)                                 | Medium       | ❌           | ✅       | ✅        |
 |         1174 | [Immediate Food Delivery II](https://leetcode.com/problems/immediate-food-delivery-ii)                         | Medium       | ❌           | ✅       | ✅        |
-|          550 | [Game Play Analysis IV](https://leetcode.com/problems/game-play-analysis-iv)                                   | Medium       | ❌           | ❌       | ✅        |
+|          550 | [Game Play Analysis IV](https://leetcode.com/problems/game-play-analysis-iv)                                   | Medium       | ❌           | ✅       | ✅        |
 ## Sorting and Grouping
 |   problem_id | title                                                                                                                              | difficulty   | DataFusion   | pandas   | PyArrow   |
 |-------------:|:-----------------------------------------------------------------------------------------------------------------------------------|:-------------|:-------------|:---------|:----------|

problems/pandas.py

@@ -1,5 +1,6 @@
 """Solutions to LeetCode problems in pandas."""
 
+from datetime import timedelta
 from decimal import ROUND_HALF_UP, Decimal
 
 import pandas as pd
@@ -87,6 +88,36 @@ def problem_197(weather: pd.DataFrame) -> pd.DataFrame:
     return weather[mask][["id"]]
 
 
+def problem_550(activity: pd.DataFrame) -> pd.DataFrame:
+    """Report the fraction of players who logged in the day after their first login.
+
+    Write a solution to report the fraction of players that logged in again on the
+    day after the day they first logged in, rounded to 2 decimal places. In other
+    words, you need to count the number of players that logged in for at least two
+    consecutive days starting from their first login date, then divide that number by
+    the total number of players.
+
+    Parameters
+    ----------
+    activity : pd.DataFrame
+        This table shows the activity of players of some games.
+
+    Returns
+    -------
+    pd.DataFrame
+
+    """
+    grouped = activity.groupby("player_id", as_index=False).aggregate(
+        min_event_date=pd.NamedAgg("event_date", "min")
+    )
+    grouped["event_date"] = grouped["min_event_date"] + timedelta(days=1)
+    joined = grouped.merge(activity, how="left", on=["player_id", "event_date"])
+    joined.device_id.notna().sum() / len(joined.index)
+    return pd.DataFrame(
+        [joined.games_played.notna().sum() / len(joined.index)], columns=["fraction"]
+    ).round(2)
+
+
 def problem_620(cinema: pd.DataFrame) -> pd.DataFrame:
     """Report movies with odd IDs and descriptions not equal to "boring".
 

tests/test_pandas.py

@@ -153,6 +153,62 @@ def test_problem_197(input_data, expected_data):
     )
 
 
+@pytest.mark.parametrize(
+    "input_data, expected_data",
+    [
+        pytest.param(
+            {
+                "player_id": [1, 1, 2, 3, 3],
+                "device_id": [2, 2, 3, 1, 4],
+                "event_date": [
+                    datetime(2016, 3, 1),
+                    datetime(2016, 3, 2),
+                    datetime(2017, 6, 25),
+                    datetime(2016, 3, 2),
+                    datetime(2018, 7, 3),
+                ],
+                "games_played": [5, 6, 1, 0, 5],
+            },
+            {"fraction": [0.33]},
+            id="happy_path_basic",
+        ),
+        pytest.param(
+            {
+                "player_id": [1, 1, 1, 2, 2],
+                "device_id": [2, 2, 3, 1, 4],
+                "event_date": [
+                    datetime(2023, 1, 1),
+                    datetime(2023, 1, 2),
+                    datetime(2023, 1, 3),
+                    datetime(2023, 1, 1),
+                    datetime(2023, 1, 2),
+                ],
+                "games_played": [1, 2, 3, 4, 5],
+            },
+            {"fraction": [1.0]},
+            id="happy_path_multiple_dates",
+        ),
+        pytest.param(
+            {
+                "player_id": [1],
+                "device_id": [1],
+                "event_date": [datetime(2023, 1, 1)],
+                "games_played": [1],
+            },
+            {"fraction": [0.0]},
+            id="edge_case_single_entry",
+        ),
+    ],
+)
+def test_problem_550(input_data, expected_data):
+    table = pd.DataFrame(input_data)
+    expected_table = pd.DataFrame(expected_data)
+    result = problem_550(table).reset_index(drop=True)
+    assert_frame_equal(
+        result, expected_table, check_dtype=False, check_index_type=False
+    )
+
+
 @pytest.mark.parametrize(
     "input_data, expected_data",
     [

tests/test_pyarrow.py

@@ -274,6 +274,7 @@ def test_problem_197(input_data, expected_data):
         pytest.param(
             {
                 "player_id": [1, 1, 1, 2, 2],
+                "device_id": [2, 2, 3, 1, 4],
                 "event_date": [
                     datetime(2023, 1, 1),
                     datetime(2023, 1, 2),
@@ -289,6 +290,7 @@ def test_problem_197(input_data, expected_data):
         pytest.param(
             {
                 "player_id": [1],
+                "device_id": [1],
                 "event_date": [datetime(2023, 1, 1)],
                 "games_played": [1],
             },