Find interview candidates

# Table: Contests
# 
# +--------------+------+
# | Column Name  | Type |
# +--------------+------+
# | contest_id   | int  |
# | gold_medal   | int  |
# | silver_medal | int  |
# | bronze_medal | int  |
# +--------------+------+
# contest_id is the column with unique values for this table.
# This table contains the LeetCode contest ID and the user IDs of the gold, silver, and bronze medalists.
# It is guaranteed that any consecutive contests have consecutive IDs and that no ID is skipped.
# 
#  
# 
# Table: Users
# 
# +-------------+---------+
# | Column Name | Type    |
# +-------------+---------+
# | user_id     | int     |
# | mail        | varchar |
# | name        | varchar |
# +-------------+---------+
# user_id is the column with unique values for this table.
# This table contains information about the users.
# 
#  
# 
# Write a solution to report the name and the mail of all interview candidates. A user is an interview candidate if at least one of these two conditions is true:
# 
# The user won any medal in three or more consecutive contests.
# 
# The user won the gold medal in three or more different contests (not necessarily consecutive).
# 
# Return the result table in any order.
# 
# The result format is in the following example.
#
# Example 1:
# Input:
# Contests table:
# +------------+------------+--------------+--------------+
# | contest_id | gold_medal | silver_medal | bronze_medal |
# +------------+------------+--------------+--------------+
# | 190        | 1          | 5            | 2            |
# | 191        | 2          | 3            | 5            |
# | 192        | 5          | 2            | 3            |
# | 193        | 1          | 3            | 5            |
# | 194        | 4          | 5            | 2            |
# | 195        | 4          | 2            | 1            |
# | 196        | 1          | 5            | 2            |
# +------------+------------+--------------+--------------+
# Users table:
# +---------+--------------------+-------+
# | user_id | mail               | name  |
# +---------+--------------------+-------+
# | 1       | sarah@leetcode.com | Sarah |
# | 2       | bob@leetcode.com   | Bob   |
# | 3       | alice@leetcode.com | Alice |
# | 4       | hercy@leetcode.com | Hercy |
# | 5       | quarz@leetcode.com | Quarz |
# +---------+--------------------+-------+
# Output:
# +-------+--------------------+
# | name  | mail               |
# +-------+--------------------+
# | Sarah | sarah@leetcode.com |
# | Bob   | bob@leetcode.com   |
# | Alice | alice@leetcode.com |
# | Quarz | quarz@leetcode.com |
# +-------+--------------------+
# Explanation:
# Sarah won 3 gold medals (190, 193, and 196), so we include her in the result table.
# Bob won a medal in 3 consecutive contests (190, 191, and 192), so we include him in the result table.
# - Note that he also won a medal in 3 other consecutive contests (194, 195, and 196).
# Alice won a medal in 3 consecutive contests (191, 192, and 193), so we include her in the result table.
# Quarz won a medal in 5 consecutive contests (190, 191, 192, 193, and 194), so we include them in the result table.

import pandas as pd

def find_interview_candidates(contests: pd.DataFrame, users: pd.DataFrame) -> pd.DataFrame:
  # Condition 1: gold medal in 3+ different contests
  gold_counts = contests.groupby('gold_medal').size().reset_index(name='count')
  gold_candidates = set(gold_counts[gold_counts['count'] >= 3]['gold_medal'])

  # Condition 2: any medal in 3+ consecutive contests
  contests = contests.sort_values('contest_id')
  melted = contests.melt(id_vars='contest_id', value_vars=['gold_medal', 'silver_medal', 'bronze_medal'],
               value_name='user_id')
  melted = melted[['contest_id', 'user_id']].drop_duplicates().sort_values(['user_id', 'contest_id'])

  consecutive_candidates = set()
  for user_id, group in melted.groupby('user_id'):
    contest_ids = sorted(group['contest_id'].values)
    streak = 1
    for i in range(1, len(contest_ids)):
      if contest_ids[i] == contest_ids[i - 1] + 1:
        streak += 1
        if streak >= 3:
          consecutive_candidates.add(user_id)
          break
      else:
        streak = 1

  all_candidates = gold_candidates | consecutive_candidates
  result = users[users['user_id'].isin(all_candidates)][['name', 'mail']]
  return result