Arrange Table by Gender — LeetCode 2308 Python Solution

Table: Genders +-------------+---------+ | Column Name | Type | +-------------+---------+ | user_id | int | | gender | varchar | +-------------+---------+ user_id is the primary key (column with unique values) for this table. gender is ENUM (category) of type 'female', 'male', or 'other'.

+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| user_id     | int     |
| gender      | varchar |
+-------------+---------+
user_id is the primary key (column with unique values) for this table.
gender is ENUM (category) of type 'female', 'male', or 'other'.
Each row in this table contains the ID of a user and their gender.
The table has an equal number of 'female', 'male', and 'other'.

import duckdb
import pandas as pd

def solution(genders: pd.DataFrame) -> pd.DataFrame:
    con = duckdb.connect()
    con.register("Genders", genders)
    return con.execute("""WITH
    t AS (
        SELECT
            *,
            RANK() OVER (
                PARTITION BY gender
                ORDER BY user_id
            ) AS rk1,
            CASE
                WHEN gender = 'female' THEN 0
                WHEN gender = 'other' THEN 1
                ELSE 2
            END AS rk2
        FROM Genders
    )
SELECT user_id, gender
FROM t
ORDER BY rk1, rk2;""").df()