595. Big Countries

Problem¶

Given a World table, return name, population, and area for countries where area >= 3000000 or population >= 25000000.

Input:

# World
----------+-----------+---------+------------+----
name      | continent | area    | population | gdp
----------+-----------+---------+------------+----
Small     | X         | 100     | 10         | 1
AreaOnly  | X         | 3000000 | 10         | 1
PopOnly   | X         | 100     | 25000000   | 1
NotBig    | X         | 2999999 | 24999999   | 1

Output:

---------+------------+--------
name     | population | area
---------+------------+--------
AreaOnly | 10         | 3000000
PopOnly  | 25000000   | 100

Key trick¶

Use OR, not AND, because either condition is enough for a country to be big.

Trap¶

Using AND instead of OR.
Using > instead of >=, which incorrectly excludes boundary values.
Returning extra columns like continent or gdp.
Assuming output order matters.

Why is it interesting?¶

This is a simple filtering problem that checks whether you can translate business rules exactly into boolean conditions.

SQL solution¶

-- SQLite
-- Select only the required columns.
-- A country is big if either threshold is met.
SELECT
    name,
    population,
    area
FROM World
WHERE area >= 3000000
   OR population >= 25000000;

Pandas solution¶

import pandas as pd

def big_countries(world: pd.DataFrame) -> pd.DataFrame:
    # Parentheses are required because | has higher precedence issues with comparisons.
    mask = (world["area"] >= 3_000_000) | (world["population"] >= 25_000_000)

    # Return only the requested columns.
    return world.loc[mask, ["name", "population", "area"]]

Pytest test¶

import sqlite3

import pandas as pd
import pytest


SQL_QUERY = """
SELECT
    name,
    population,
    area
FROM World
WHERE area >= 3000000
   OR population >= 25000000;
"""


def big_countries(world: pd.DataFrame) -> pd.DataFrame:
    mask = (world["area"] >= 3_000_000) | (world["population"] >= 25_000_000)
    return world.loc[mask, ["name", "population", "area"]]


def run_sql(rows):
    conn = sqlite3.connect(":memory:")
    conn.execute(
        """
        CREATE TABLE World (
            name TEXT PRIMARY KEY,
            continent TEXT,
            area INTEGER,
            population INTEGER,
            gdp INTEGER
        );
        """
    )
    conn.executemany(
        """
        INSERT INTO World (name, continent, area, population, gdp)
        VALUES (?, ?, ?, ?, ?);
        """,
        rows,
    )

    result = pd.read_sql_query(SQL_QUERY, conn)
    conn.close()
    return result


def normalize(df: pd.DataFrame) -> pd.DataFrame:
    return (
        df[["name", "population", "area"]]
        .sort_values(["name", "population", "area"])
        .reset_index(drop=True)
    )


@pytest.mark.parametrize(
    "rows, expected_rows",
    [
        (
            [
                ("Afghanistan", "Asia", 652230, 25500100, 20343000000),
                ("Albania", "Europe", 28748, 2831741, 12960000000),
                ("Algeria", "Africa", 2381741, 37100000, 188681000000),
                ("Andorra", "Europe", 468, 78115, 3712000000),
                ("Angola", "Africa", 1246700, 20609294, 100990000000),
            ],
            [
                ("Afghanistan", 25500100, 652230),
                ("Algeria", 37100000, 2381741),
            ],
        ),
        (
            [
                ("AreaBoundary", "X", 3000000, 1, 1),
                ("PopBoundary", "X", 1, 25000000, 1),
                ("BelowBoth", "X", 2999999, 24999999, 1),
            ],
            [
                ("AreaBoundary", 1, 3000000),
                ("PopBoundary", 25000000, 1),
            ],
        ),
        (
            [
                ("BothBig", "X", 4000000, 30000000, 1),
                ("AreaOnly", "X", 4000000, 10, 1),
                ("PopOnly", "X", 10, 30000000, 1),
                ("Small", "X", 10, 10, 1),
            ],
            [
                ("BothBig", 30000000, 4000000),
                ("AreaOnly", 10, 4000000),
                ("PopOnly", 30000000, 10),
            ],
        ),
        (
            [
                ("SmallA", "X", 1, 1, 1),
                ("SmallB", "X", 2999999, 24999999, 1),
            ],
            [],
        ),
    ],
)
def test_big_countries_sql_and_pandas(rows, expected_rows):
    world = pd.DataFrame(
        rows,
        columns=["name", "continent", "area", "population", "gdp"],
    )

    expected = pd.DataFrame(
        expected_rows,
        columns=["name", "population", "area"],
    )

    sql_result = run_sql(rows)
    pandas_result = big_countries(world)

    pd.testing.assert_frame_equal(
        normalize(sql_result),
        normalize(expected),
        check_dtype=False,
    )
    pd.testing.assert_frame_equal(
        normalize(pandas_result),
        normalize(expected),
        check_dtype=False,
    )

Comment on my solution¶

Your SQL and Pandas solutions are correct and interview-ready.

The boolean condition uses OR, which matches the requirement.
The thresholds correctly use >=, so boundary values are included.
The returned columns are exactly the requested ones.
The Pandas expression correctly uses parentheses around each comparison.

SELECT
    name,
    population,
    area
FROM World
WHERE area >= 3000000 OR population >= 25000000;

import pandas as pd

def big_countries(world: pd.DataFrame) -> pd.DataFrame:
    return world.loc[(world["area"] >= 3000000) | (world["population"] >= 25000000),
                     ["name", "population", "area"]]