[Python/Pandas] 리트코드 1050. Actors and Directors Who Cooperated At Least Three Times

문제

Table: ActorDirector
+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| actor_id    | int     |
| director_id | int     |
| timestamp   | int     |
+-------------+---------+

timestamp is the primary key (column with unique values) for this table.

Write a solution to find all the pairs (actor_id, director_id) where the actor has cooperated with the director at least three times.
Return the result table in any order.

 

https://leetcode.com/problems/actors-and-directors-who-cooperated-at-least-three-times/description/?lang=pythondata

예시

Input: 
ActorDirector table:
+-------------+-------------+-------------+
| actor_id    | director_id | timestamp   |
+-------------+-------------+-------------+
| 1           | 1           | 0           |
| 1           | 1           | 1           |
| 1           | 1           | 2           |
| 1           | 2           | 3           |
| 1           | 2           | 4           |
| 2           | 1           | 5           |
| 2           | 1           | 6           |
+-------------+-------------+-------------+
Output: 
+-------------+-------------+
| actor_id    | director_id |
+-------------+-------------+
| 1           | 1           |
+-------------+-------------+

Explanation: The only pair is (1, 1) where they cooperated exactly 3 times.

문제 풀이

import pandas as pd

def actors_and_directors(actor_director: pd.DataFrame) -> pd.DataFrame:
    pair_count = actor_director.groupby(['actor_id','director_id'])['timestamp'].count().reset_index()
    result = pair_count.loc[pair_count['timestamp'] >= 3, ['actor_id', 'director_id']]
    return result

---

파이썬을 독학하시는 분들에게 도움이 되길 바라며,

혹 더 좋은 방법이 있거나 오류가 있다면 편하게 말씀 부탁드립니다.