This video provides a walkthrough of the challenge requirements for segmenting a dataset by two dimensions (such as acquisition month and user region) and calculating advanced retention metrics. You will see how to use pandas to group data, create multi-level cohorts, and compute retention and churn rates for each segment.

To complete this challenge, follow these steps:

- Use a **pandas DataFrame** containing user activity data, with columns such as `user_id`, `acquisition_month`, `region`, and `activity_month`;
- Segment the users by both `acquisition_month` and `region` to create multi-level cohorts;
- For each cohort, calculate the number of users retained in each subsequent month after acquisition;
- Compute the **retention rate** for each cohort as the percentage of users active in a given month compared to the original cohort size;
- Calculate the **churn rate** as 1 minus the retention rate for each period.

You will need to use **pandas** grouping and aggregation methods to perform these calculations efficiently.

import unittest
import user_code
import importlib
import pandas as pd

class TestTask(unittest.TestCase):
    def setUp(self):
        self.data = {
            'user_id': [1, 1, 2, 2, 3, 4, 4, 5, 6, 6],
            'acquisition_month': [1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
            'activity_month': [1, 2, 1, 2, 2, 2, 3, 2, 2, 3]
        }
        self.df = pd.DataFrame(self.data)

    def test_returns_dataframe_and_columns(self):
        import user_code
        importlib.reload(user_code)
        result = user_code.calculate_retention_by_acquisition(self.df)
        expected_columns = ['cohort', 'months_since_acquisition', 'retention_rate']
        _dynamic_test(
            self,
            isinstance(result, pd.DataFrame),
            "Returned object is a pandas DataFrame.",
            f"Expected DataFrame, got {type(result)}."
        )
        _dynamic_test(
            self,
            all(col in result.columns for col in expected_columns),
            f"Result contains all required columns: {expected_columns}.",
            f"Missing columns in result. Got columns: {list(result.columns)}."
        )

    def test_cohort_definition(self):
        import user_code
        importlib.reload(user_code)
        result = user_code.calculate_retention_by_acquisition(self.df)
        cohort_values = set(result['cohort'])
        expected_cohorts = set([1, 2])
        _dynamic_test(
            self,
            cohort_values == expected_cohorts,
            f"Cohorts are defined as acquisition_month: {expected_cohorts}",
            f"Expected cohorts: {expected_cohorts}, got: {cohort_values}"
        )

    def test_retention_rate_calculation(self):
        import user_code
        importlib.reload(user_code)
        result = user_code.calculate_retention_by_acquisition(self.df)
        # For cohort 1, months_since_acquisition=0: users 1 and 2
        # cohort size: users 1 and 2, so 2
        # active users at month 0: 2
        # retention rate: 2/2 = 1.0
        row = result[(result['cohort'] == 1) & (result['months_since_acquisition'] == 0)]
        rr = float(row['retention_rate'].iloc[0]) if not row.empty else None
        _dynamic_test(
            self,
            rr is not None and abs(rr - 1.0) < 1e-8,
            "Retention rate for cohort 1 at month 0 is correct (1.0)",
            f"Expected 1.0, got {rr}"
        )
        # For cohort 1, months_since_acquisition=1: users 1 and 2 again, so 2/2 = 1.0
        row = result[(result['cohort'] == 1) & (result['months_since_acquisition'] == 1)]
        rr = float(row['retention_rate'].iloc[0]) if not row.empty else None
        _dynamic_test(
            self,
            rr is not None and abs(rr - 1.0) < 1e-8,
            "Retention rate for cohort 1 at month 1 is correct (1.0)",
            f"Expected 1.0, got {rr}"
        )
        # For cohort 2, months_since_acquisition=0: users 3, 4, 5, 6
        # cohort size: users 3, 4, 5, 6, so 4
        # active users at month 0: 4
        # retention rate: 4/4 = 1.0
        row = result[(result['cohort'] == 2) & (result['months_since_acquisition'] == 0)]
        rr = float(row['retention_rate'].iloc[0]) if not row.empty else None
        _dynamic_test(
            self,
            rr is not None and abs(rr - 1.0) < 1e-8,
            "Retention rate for cohort 2 at month 0 is correct (1.0)",
            f"Expected 1.0, got {rr}"
        )
        # For cohort 2, months_since_acquisition=1: users 4, 6, so 2/4 = 0.5
        row = result[(result['cohort'] == 2) & (result['months_since_acquisition'] == 1)]
        rr = float(row['retention_rate'].iloc[0]) if not row.empty else None
        _dynamic_test(
            self,
            rr is not None and abs(rr - 0.5) < 1e-8,
            "Retention rate for cohort 2 at month 1 is correct (0.5)",
            f"Expected 0.5, got {rr}"
        )

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)


Challenge: Advanced Segmentation and Retention

Solution