585. Investments in 2016

Problem¶

Given Insurance, return the rounded sum of tiv_2016 for policies where:

tiv_2015 appears in at least two rows.
The location pair (lat, lon) appears in exactly one row.

Input:

# Insurance
----+----------+----------+-----+----
pid | tiv_2015 | tiv_2016 | lat | lon
----+----------+----------+-----+----
1   | 10       | 5        | 10  | 10
2   | 20       | 20       | 20  | 20
3   | 10       | 30       | 20  | 20
4   | 10       | 40       | 40  | 40

Output:

tiv_2016
--------
45.00

Key trick¶

Use group counts as row-level filters:

Count rows per tiv_2015; keep rows with count greater than 1.
Count rows per (lat, lon); keep rows with count equal to 1.

Trap¶

Checking lat and lon uniqueness separately instead of checking the pair (lat, lon).
Rounding each tiv_2016 before summing instead of rounding the final sum.
Joining grouped tables incorrectly and duplicating rows.
In PostgreSQL, ROUND(double precision, integer) is not valid, so the sum must be cast to numeric.

Why is it interesting?¶

It tests whether you can combine aggregate conditions with row-level filtering, which is exactly what window functions or groupby().transform() are good at.

SQL solution¶

SQLite¶

WITH marked AS (
    SELECT
        tiv_2016,

        -- Same 2015 investment value as at least one other policy.
        COUNT(*) OVER (
            PARTITION BY tiv_2015
        ) AS tiv_2015_count,

        -- Unique city means the full coordinate pair appears once.
        COUNT(*) OVER (
            PARTITION BY lat, lon
        ) AS location_count
    FROM Insurance
)
SELECT
    -- Round only the final sum.
    ROUND(SUM(tiv_2016), 2) AS tiv_2016
FROM marked
WHERE tiv_2015_count > 1
  AND location_count = 1;

PostgreSQL¶

WITH marked AS (
    SELECT
        tiv_2016,

        -- Same 2015 investment value as at least one other policy.
        COUNT(*) OVER (
            PARTITION BY tiv_2015
        ) AS tiv_2015_count,

        -- Unique city means the full coordinate pair appears once.
        COUNT(*) OVER (
            PARTITION BY lat, lon
        ) AS location_count
    FROM Insurance
)
SELECT
    -- PostgreSQL needs numeric for ROUND(value, decimals).
    ROUND(SUM(tiv_2016)::numeric, 2) AS tiv_2016
FROM marked
WHERE tiv_2015_count > 1
  AND location_count = 1;

Pandas solution¶

import pandas as pd


def find_investments(insurance: pd.DataFrame) -> pd.DataFrame:
    marked = insurance.assign(
        # Broadcast each tiv_2015 group size back to its original rows.
        tiv_2015_count=insurance.groupby("tiv_2015")["pid"].transform("size"),

        # Broadcast each location-pair group size back to its original rows.
        location_count=insurance.groupby(["lat", "lon"])["pid"].transform("size"),
    )

    valid_rows = marked[
        (marked["tiv_2015_count"] > 1)
        & (marked["location_count"] == 1)
    ]

    # min_count=1 matches SQL SUM behavior when no rows qualify.
    total = valid_rows["tiv_2016"].sum(min_count=1)

    if pd.notna(total):
        total = round(float(total), 2)
    else:
        total = pd.NA

    return pd.DataFrame({"tiv_2016": [total]})

data = [[1, 10, 5, 10, 10], [2, 20, 20, 20, 20], [3, 10, 30, 20, 20], [4, 10, 40, 40, 40]]
insurance = pd.DataFrame(data, columns=['pid', 'tiv_2015', 'tiv_2016', 'lat', 'lon']).astype({'pid':'Int64', 'tiv_2015':'Float64', 'tiv_2016':'Float64', 'lat':'Float64', 'lon':'Float64'})

insurance
#    pid  tiv_2015  tiv_2016   lat   lon
# 0    1      10.0       5.0  10.0  10.0
# 1    2      20.0      20.0  20.0  20.0
# 2    3      10.0      30.0  20.0  20.0
# 3    4      10.0      40.0  40.0  40.0

# Broadcast each tiv_2015 group size back to its original rows.
insurance.assign(
        tiv_2015_count=insurance.groupby("tiv_2015")["pid"].transform("size")
)
#    pid  tiv_2015  tiv_2016   lat   lon  tiv_2015_count
# 0    1      10.0       5.0  10.0  10.0               3
# 1    2      20.0      20.0  20.0  20.0               1
# 2    3      10.0      30.0  20.0  20.0               3
# 3    4      10.0      40.0  40.0  40.0               3

insurance.groupby("tiv_2015")["pid"].transform("size")
# 0    3
# 1    1
# 2    3
# 3    3
# Name: pid, dtype: Int64

insurance.groupby("tiv_2015")["pid"].size()
# tiv_2015
# 10.0    3
# 20.0    1
# Name: pid, dtype: Int64

Pytest test¶

import sqlite3

import pandas as pd
import pytest


SQLITE_QUERY = """
WITH marked AS (
    SELECT
        tiv_2016,
        COUNT(*) OVER (
            PARTITION BY tiv_2015
        ) AS tiv_2015_count,
        COUNT(*) OVER (
            PARTITION BY lat, lon
        ) AS location_count
    FROM Insurance
)
SELECT
    ROUND(SUM(tiv_2016), 2) AS tiv_2016
FROM marked
WHERE tiv_2015_count > 1
  AND location_count = 1;
"""


def pandas_solution(insurance: pd.DataFrame) -> pd.DataFrame:
    marked = insurance.assign(
        tiv_2015_count=insurance.groupby("tiv_2015")["pid"].transform("size"),
        location_count=insurance.groupby(["lat", "lon"])["pid"].transform("size"),
    )

    valid_rows = marked[
        (marked["tiv_2015_count"] > 1)
        & (marked["location_count"] == 1)
    ]

    total = valid_rows["tiv_2016"].sum(min_count=1)

    if pd.notna(total):
        total = round(float(total), 2)
    else:
        total = pd.NA

    return pd.DataFrame({"tiv_2016": [total]})


@pytest.mark.parametrize(
    "rows, expected",
    [
        (
            [
                [1, 10.0, 5.0, 10.0, 10.0],
                [2, 20.0, 20.0, 20.0, 20.0],
                [3, 10.0, 30.0, 20.0, 20.0],
                [4, 10.0, 40.0, 40.0, 40.0],
            ],
            45.00,
        ),
        (
            [
                [1, 1.0, 1.111, 0.0, 0.0],
                [2, 1.0, 2.225, 1.0, 1.0],
            ],
            3.34,
        ),
        (
            [
                [1, 10.0, 5.0, 0.0, 0.0],
                [2, 10.0, 6.0, 0.0, 0.0],
                [3, 20.0, 7.0, 2.0, 2.0],
                [4, 30.0, 8.0, 3.0, 3.0],
                [5, 30.0, 9.0, 4.0, 4.0],
            ],
            17.00,
        ),
        (
            [
                [1, 10.0, 5.0, 1.0, 1.0],
                [2, 20.0, 6.0, 2.0, 2.0],
            ],
            None,
        ),
    ],
)
def test_investments_sqlite_and_pandas(rows, expected):
    columns = ["pid", "tiv_2015", "tiv_2016", "lat", "lon"]
    insurance = pd.DataFrame(rows, columns=columns)

    with sqlite3.connect(":memory:") as connection:
        insurance.to_sql(
            "Insurance",
            connection,
            index=False,
            if_exists="replace",
        )
        sql_value = pd.read_sql_query(SQLITE_QUERY, connection).iloc[0, 0]

    pandas_value = pandas_solution(insurance).iloc[0, 0]

    if expected is None:
        assert pd.isna(sql_value)
        assert pd.isna(pandas_value)
    else:
        assert round(float(sql_value), 2) == pytest.approx(expected)
        assert round(float(pandas_value), 2) == pytest.approx(expected)

Comment on my solution¶

Your working SQL solution is correct, but it is PostgreSQL-specific because of the ::numeric cast.

Your working Pandas solution is correct, but the merge can be avoided with groupby().transform("size"), which is usually cleaner for row-level filters based on group counts.

The failing Pandas version breaks because isin() with a MultiIndex does not compare row pairs from insurance[["lat", "lon"]] the way you intend; use a merge, a tuple key, or groupby().transform() instead.

-- WORKS
WITH valid_pids AS (
    SELECT pid
    FROM (
        SELECT
           pid,
           COUNT(*) OVER (PARTITION BY tiv_2015) AS cnt_tiv_2015,
           COUNT(*) OVER (PARTITION BY lat, lon) AS cnt_lat_lon
        FROM Insurance
    )
    WHERE cnt_tiv_2015 > 1 AND cnt_lat_lon = 1
)
SELECT
    ROUND(SUM(tiv_2016)::numeric, 2) AS tiv_2016
FROM Insurance AS i
JOIN valid_pids AS vp
  ON i.pid = vp.pid;

import pandas as pd

# WORKS
def find_investments(insurance: pd.DataFrame) -> pd.DataFrame:
    valid_tiv_2015 = insurance.groupby("tiv_2015").size()
    valid_tiv_2015 = valid_tiv_2015[valid_tiv_2015 > 1].index

    valid_lat_lon = (
        insurance.groupby(["lat", "lon"], as_index=False)
        .size()
        .loc[lambda df: df["size"] == 1]
        [["lat", "lon"]]
    )

    valid_pids = insurance.loc[insurance["tiv_2015"].isin(valid_tiv_2015)]
    valid_pids = valid_pids.merge(
        valid_lat_lon,
        on=["lat", "lon"],
        how="inner"
    )

    tiv_2016 = valid_pids["tiv_2016"].sum().round(2)
    return pd.DataFrame({"tiv_2016": [tiv_2016]})



# Error: Cannot index with multidimensional key
def find_investments(insurance: pd.DataFrame) -> pd.DataFrame:
    valid_tiv_2015 = insurance.groupby("tiv_2015").size()
    valid_tiv_2015 = valid_tiv_2015[valid_tiv_2015 > 1].index
    valid_lat_lon = insurance.groupby(["lat", "lon"]).size()
    valid_lat_lon = valid_lat_lon[valid_lat_lon == 1].index
    valid_pids = insurance.loc[insurance[["lat", "lon"]].isin(valid_lat_lon)]
    return valid_pids

Extra¶

Complement on Pandas MultiIndex and other solutions¶

In the Pandas comments, what do you mean by a tuple key? Is it possible on MultiIndex?

tuple key¶

A tuple key means turning the two columns into one comparable value:

(lat, lon)

So instead of asking:

Is lat valid?
Is lon valid?

You ask:

Is (lat, lon) valid?

That matters because location uniqueness is about the pair, not each column separately.

Why your failing code fails¶

You had roughly:

valid_lat_lon = insurance.groupby(["lat", "lon"]).size()
valid_lat_lon = valid_lat_lon[valid_lat_lon == 1].index

valid_pids = insurance.loc[
    insurance[["lat", "lon"]].isin(valid_lat_lon)
]

Problem:

insurance[["lat", "lon"]] is a DataFrame.
.isin(valid_lat_lon) does elementwise checks.
It does not compare row pairs like (lat, lon).
It returns a 2D boolean DataFrame.
.loc[...] expects a 1D row mask, so you get:

Cannot index with multidimensional key

Option 1: use tuple keys¶

location_key = pd.Series(
    list(zip(insurance["lat"], insurance["lon"])),
    index=insurance.index,
)

location_counts = location_key.value_counts()

valid_locations = location_counts[location_counts == 1].index

valid_rows = insurance[location_key.isin(valid_locations)]

Here location_key looks like:

0    (10.0, 10.0)
1    (20.0, 20.0)
2    (20.0, 20.0)
3    (40.0, 40.0)
dtype: object

So .isin(...) compares full pairs.

Option 2: use `MultiIndex`¶

Yes, this is possible and often cleaner.

location_counts = insurance.groupby(["lat", "lon"]).size()

valid_locations = location_counts[location_counts == 1].index

row_locations = pd.MultiIndex.from_frame(insurance[["lat", "lon"]])

valid_rows = insurance[row_locations.isin(valid_locations)]

This works because:

valid_locations is a MultiIndex.
row_locations is also a MultiIndex.
.isin() compares full index keys like (lat, lon).

Full Pandas solution using `MultiIndex`¶

import pandas as pd


def find_investments(insurance: pd.DataFrame) -> pd.DataFrame:
    tiv_2015_counts = insurance.groupby("tiv_2015").size()
    valid_tiv_2015 = tiv_2015_counts[tiv_2015_counts > 1].index

    location_counts = insurance.groupby(["lat", "lon"]).size()
    valid_locations = location_counts[location_counts == 1].index

    row_locations = pd.MultiIndex.from_frame(insurance[["lat", "lon"]])

    mask = (
        insurance["tiv_2015"].isin(valid_tiv_2015)
        & row_locations.isin(valid_locations)
    )

    total = insurance.loc[mask, "tiv_2016"].sum().round(2)

    return pd.DataFrame({"tiv_2016": [total]})

Best interview version¶

The cleanest Pandas answer is still usually transform, because it avoids manual joins, tuple keys, and MultiIndex handling:

import pandas as pd


def find_investments(insurance: pd.DataFrame) -> pd.DataFrame:
    tiv_2015_count = insurance.groupby("tiv_2015")["pid"].transform("size")
    location_count = insurance.groupby(["lat", "lon"])["pid"].transform("size")

    mask = (tiv_2015_count > 1) & (location_count == 1)

    total = insurance.loc[mask, "tiv_2016"].sum().round(2)

    return pd.DataFrame({"tiv_2016": [total]})

Mental model¶

Use this rule:

If you need column-by-column membership, use DataFrame .isin(...).
If you need row-pair membership, create one key:
- Tuple key: (lat, lon)
- MultiIndex: index key (lat, lon)
- merge: relational key ["lat", "lon"]

For this problem, lat and lon must be treated as one combined key.