[ 알고리즘 ] LeetCode 2153. The Number of Passengers in Each Bus II

블로그를 이전 중이라 완료되기 전까지는 벨로그에 작성할 계획입니다.
이후 모든 글은 https://weekwith.me 에 작성 예정이니 다른 글이 궁금하시다면 해당 링크를 통해 방문해주세요.

본 글은 [ LeetCode ] 2153. The Number of Passengers in Each Bus II를 풀고 작성한 글입니다.

문제

테이블

Table: Buses

+--------------+------+
| Column Name  | Type |
+--------------+------+
| bus_id       | int  |
| arrival_time | int  |
| capacity     | int  |
+--------------+------+
bus_id is the primary key column for this table.
Each row of this table contains information about the arrival time of a bus at the LeetCode station and its capacity (the number of empty seats it has).
No two buses will arrive at the same time and all bus capacities will be positive integers.

Table: Passengers

+--------------+------+
| Column Name  | Type |
+--------------+------+
| passenger_id | int  |
| arrival_time | int  |
+--------------+------+
passenger_id is the primary key column for this table.
Each row of this table contains information about the arrival time of a passenger at the LeetCode station.

요구사항

Buses and passengers arrive at the LeetCode station. If a bus arrives at the station at a time Tbus and a passenger arrived at a time Tpassenger where Tpassenger <= Tbus and the passenger did not catch any bus, the passenger will use that bus. In addition, each bus has a capacity. If at the moment the bus arrives at the station there are more passengers waiting than its capacity capacity , only capacity passengers will use the bus.

Write an SQL query to report the number of users that used each bus.

Return the result table ordered by bus_id in ascending order .

풀이

접근법

누적된 차이를 계속해서 비교해가며 구해야하기 때문에 WITH RECURSIVE 공통 테이블 표현식(CTE_Common Table Expression)을 활용해서 문제를 해결할 수 있다.

이때 유의해야 할 점은 bus_id 필드의 값이 곧 arrival_time 필드의 값이 오름차순 기준으로 정렬된 순서와 동일하다는 게 보장되지 않기 때문에 공통 테이블 표현식의 재귀 부분인 JOIN 구의 조건에 사용하기 위해 arrival_time 필드를 오름차순 기준으로 정렬해서 ROW_NUMBER 윈도우 함수(Window Function)를 통해 행 번호를 구해야 한다.

추가적으로 LAG 윈도우 함수의 두 번째 매개변수는 윈도우의 범위를 의미하며 세 번째 매개변수는 만약 그 값이 NULL 일 경우에 입력할 기본값을 의미한다. 기본값은 NULL 인데 아래 풀이에서는 -1 을 인자로 넘겼다.

누적된 차이를 구하는 부분은 생각했지만 그 방법이 JOIN 구와 함께 ROW_NUMBER 윈도우 함수를 사용하는 것이란 걸 쉽게 떠올리지 못했다. 단순히 WHERE 구 뿐만 아니라 JOIN 구를 활용해서 공통 테이블 표현식을 사용하는 경우를 잘 기억해두어야겠다.

나의 풀이

접근법을 토대로 문제를 해결하면 아래와 같다.

WITH RECURSIVE cte1 (row_num, bus_id, capacity, waited_passengers_cnt) AS (
    SELECT
        ROW_NUMBER() OVER(ORDER BY Periods.arrival_time ASC) AS row_num,
        Periods.bus_id,
        Periods.capacity,
        COUNT(Passengers.passenger_id) AS waited_passengers_cnt
    FROM (
        SELECT
            bus_id,
            capacity,            
            arrival_time,
            LAG(arrival_time, 1, -1) OVER(ORDER BY arrival_time ASC) AS previous_time
        FROM Buses
    ) AS Periods
    LEFT JOIN Passengers
    ON (
        Periods.arrival_time >= Passengers.arrival_time
        AND
        Periods.previous_time < Passengers.arrival_time
    )
    GROUP BY Periods.bus_id
), cte2 (row_num, bus_id, diff, passengers_cnt) AS (
    SELECT
        row_num,
        bus_id,
        IF(capacity < waited_passengers_cnt, waited_passengers_cnt - capacity, 0) AS diff,
        IF(capacity < waited_passengers_cnt, capacity, waited_passengers_cnt) AS passengers_cnt
    FROM cte1
    WHERE row_num = 1
    UNION ALL
    SELECT
        cte1.row_num,
        cte1.bus_id,
        IF(cte1.capacity < cte1.waited_passengers_cnt + cte2.diff, cte1.waited_passengers_cnt + cte2.diff - cte1.capacity, 0) AS diff,
        IF(cte1.capacity < cte1.waited_passengers_cnt + cte2.diff, cte1.capacity, cte1.waited_passengers_cnt + cte2.diff) AS passengers_cnt
    FROM cte1
    JOIN cte2
    ON cte1.row_num = cte2.row_num + 1
)

SELECT bus_id, passengers_cnt
FROM cte2
ORDER BY bus_id ASC;