602. Friend Requests II: Who Has the Most Friends

Problem¶

Given accepted friend pairs (requester_id, accepter_id), treat friendship as undirected and return the user with the largest number of friends and that count.

Example:

Input

# RequestAccepted
-------------+-------------+------------
requester_id | accepter_id | accept_date
-------------+-------------+------------
1            | 2           | 2016-06-03
1            | 3           | 2016-06-08
2            | 3           | 2016-06-08
3            | 4           | 2016-06-09

Output

---+----
id | num
---+----
3  | 3

Key trick¶

Turn each accepted request into two rows, one for each endpoint, then count friends per user.

Trap¶

Counting only requester_id or only accepter_id.
Forgetting friendship is undirected.
Using UNION instead of UNION ALL.
- UNION removes duplicate rows and can undercount.
Overthinking with self-joins when simple aggregation is enough.

Why is it interesting?¶

It tests whether you can convert directed-looking data into an undirected count with a clean aggregation pattern.

SQL solution¶

-- Duplicate each friendship once per endpoint, then count per person.
WITH all_friends AS (
    SELECT requester_id AS id FROM RequestAccepted
    UNION ALL
    SELECT accepter_id AS id FROM RequestAccepted
)
SELECT
    id,
    COUNT(*) AS num
FROM all_friends
GROUP BY id
ORDER BY num DESC
LIMIT 1;

Pandas solution¶

import pandas as pd

def most_friends(request_accepted: pd.DataFrame) -> pd.DataFrame:
    # Stack both endpoints into one "id" column so each friendship
    # contributes once to each person.
    all_friends = pd.concat(
        [
            request_accepted[["requester_id"]].rename(columns={"requester_id": "id"}),
            request_accepted[["accepter_id"]].rename(columns={"accepter_id": "id"}),
        ],
        ignore_index=True,
    )

    # Count occurrences per person, sort by count descending,
    # then pick the unique winner guaranteed by the prompt.
    return (
        all_friends.groupby("id", as_index=False)
        .size()
        .rename(columns={"size": "num"})
        .sort_values("num", ascending=False)
        .head(1)
        .reset_index(drop=True)
    )

Pytest test¶

import sqlite3
import pandas as pd
import pandas.testing as pdt
import pytest


SQLITE_QUERY = """
WITH all_friends AS (
    SELECT requester_id AS id FROM RequestAccepted
    UNION ALL
    SELECT accepter_id AS id FROM RequestAccepted
)
SELECT
    id,
    COUNT(*) AS num
FROM all_friends
GROUP BY id
ORDER BY num DESC
LIMIT 1;
"""


def most_friends(request_accepted: pd.DataFrame) -> pd.DataFrame:
    # Stack both endpoints into one "id" column so each friendship
    # contributes once to each person.
    all_friends = pd.concat(
        [
            request_accepted[["requester_id"]].rename(columns={"requester_id": "id"}),
            request_accepted[["accepter_id"]].rename(columns={"accepter_id": "id"}),
        ],
        ignore_index=True,
    )

    # Count occurrences per person, sort by count descending,
    # then pick the unique winner guaranteed by the prompt.
    return (
        all_friends.groupby("id", as_index=False)
        .size()
        .rename(columns={"size": "num"})
        .sort_values("num", ascending=False)
        .head(1)
        .reset_index(drop=True)
    )


@pytest.mark.parametrize(
    "rows, expected",
    [
        (
            [
                [1, 2, "2016-06-03"],
                [1, 3, "2016-06-08"],
                [2, 3, "2016-06-08"],
                [3, 4, "2016-06-09"],
            ],
            pd.DataFrame({"id": [3], "num": [3]}),
        ),
        (
            [
                [10, 20, "2020-01-01"],
            ],
            pd.DataFrame({"id": [10], "num": [1]}),
        ),
        (
            [
                [1, 5, "2020-01-01"],
                [2, 5, "2020-01-02"],
                [3, 5, "2020-01-03"],
                [4, 5, "2020-01-04"],
                [6, 5, "2020-01-05"],
            ],
            pd.DataFrame({"id": [5], "num": [5]}),
        ),
    ],
)
def test_most_friends_sql_and_pandas(rows, expected):
    df = pd.DataFrame(
        rows,
        columns=["requester_id", "accepter_id", "accept_date"],
    )
    df["requester_id"] = df["requester_id"].astype("int64")
    df["accepter_id"] = df["accepter_id"].astype("int64")
    df["accept_date"] = pd.to_datetime(df["accept_date"])

    # Test Pandas solution.
    pandas_result = most_friends(df)
    pdt.assert_frame_equal(
        pandas_result.reset_index(drop=True),
        expected.reset_index(drop=True),
    )

    # Test SQLite solution.
    conn = sqlite3.connect(":memory:")
    try:
        df.to_sql("RequestAccepted", conn, index=False, if_exists="replace")
        sql_result = pd.read_sql_query(SQLITE_QUERY, conn).astype(
            {"id": "int64", "num": "int64"}
        )
        pdt.assert_frame_equal(
            sql_result.reset_index(drop=True),
            expected.reset_index(drop=True),
        )
    finally:
        conn.close()

Comment on my solution¶

Your solution is correct and interview-good.

The main simplification is to skip the two pre-aggregations.
Just UNION ALL the two id columns first, then do one final GROUP BY.
In Pandas, the same simplification makes the code shorter and easier to explain.
Adding a secondary sort by id is a nice deterministic tie-breaker, even if the prompt guarantees one winner.

-- WORKS
WITH grouped_by_requester AS (
    SELECT
        requester_id AS id,
        COUNT(*) AS num
    FROM RequestAccepted
    GROUP BY requester_id
),
grouped_by_accepter AS (
    SELECT
        accepter_id AS id,
        COUNT(*) AS num
    FROM RequestAccepted
    GROUP BY accepter_id
)
SELECT
    id,
    SUM(num) AS num
FROM (
    SELECT * FROM grouped_by_requester
    UNION ALL
    SELECT * FROM grouped_by_accepter
)
GROUP BY id
ORDER BY num DESC
LIMIT 1;

# WORKS
import pandas as pd

def most_friends(request_accepted: pd.DataFrame) -> pd.DataFrame:
    grouped_by_requester = (
        request_accepted.groupby("requester_id", as_index=False)
        .size()
        .rename(columns={"requester_id": "id", "size": "num"})
    )
    grouped_by_accepter = (
        request_accepted.groupby("accepter_id", as_index=False)
        .size()
        .rename(columns={"accepter_id": "id", "size": "num"})
    )
    result = pd.concat([grouped_by_requester, grouped_by_accepter])
    return (
        result.groupby("id", as_index=False)
        .agg(num=("num", "sum"))
        .sort_values("num", ascending=False)
        .head(1)
    )