Find top scoring students ii

# Table: students
# 
# +-------------+----------+
# | Column Name | Type     |
# +-------------+----------+
# | student_id  | int      |
# | name        | varchar  |
# | major       | varchar  |
# +-------------+----------+
# student_id is the primary key for this table.
# Each row contains the student ID, student name, and their major.
# 
# Table: courses
# 
# +-------------+-------------------+
# | Column Name | Type              |
# +-------------+-------------------+
# | course_id   | int               |
# | name        | varchar           |
# | credits     | int               |
# | major       | varchar           |
# | mandatory   | enum              |
# +-------------+-------------------+
# course_id is the primary key for this table.
# mandatory is an enum type of ('Yes', 'No').
# Each row contains the course ID, course name, credits, major it belongs to, and whether the course is mandatory.
# 
# Table: enrollments
# 
# +-------------+----------+
# | Column Name | Type     |
# +-------------+----------+
# | student_id  | int      |
# | course_id   | int      |
# | semester    | varchar  |
# | grade       | varchar  |
# | GPA         | decimal  |
# +-------------+----------+
# (student_id, course_id, semester) is the primary key (combination of columns with unique values) for this table.
# Each row contains the student ID, course ID, semester, and grade received.
# 
# Write a solution to find the students who meet the following criteria:
# 
# Have taken all mandatory courses and at least two elective courses offered in their major.
# 
# Achieved a grade of A in all mandatory courses and at least B in elective courses.
# 
# Maintained an average GPA of at least 2.5 across all their courses (including those outside their major).
# 
# Return the result table ordered by student_id in ascending order.
#
# Example 1:
# Input:
# students table:
# +------------+------------------+------------------+
# | student_id | name             | major            |
# +------------+------------------+------------------+
# | 1          | Alice            | Computer Science |
# | 2          | Bob              | Computer Science |
# | 3          | Charlie          | Mathematics      |
# | 4          | David            | Mathematics      |
# +------------+------------------+------------------+
# courses table:
# +-----------+-------------------+---------+------------------+----------+
# | course_id | name              | credits | major            | mandatory|
# +-----------+-------------------+---------+------------------+----------+
# | 101       | Algorithms        | 3       | Computer Science | yes      |
# | 102       | Data Structures   | 3       | Computer Science | yes      |
# | 103       | Calculus          | 4       | Mathematics      | yes      |
# | 104       | Linear Algebra    | 4       | Mathematics      | yes      |
# | 105       | Machine Learning  | 3       | Computer Science | no       |
# | 106       | Probability       | 3       | Mathematics      | no       |
# | 107       | Operating Systems | 3       | Computer Science | no       |
# | 108       | Statistics        | 3       | Mathematics      | no       |
# +-----------+-------------------+---------+------------------+----------+
# enrollments table:
# +------------+-----------+-------------+-------+-----+
# | student_id | course_id | semester    | grade | GPA |
# +------------+-----------+-------------+-------+-----+
# | 1          | 101       | Fall 2023   | A     | 4.0 |
# | 1          | 102       | Spring 2023 | A     | 4.0 |
# | 1          | 105       | Spring 2023 | A     | 4.0 |
# | 1          | 107       | Fall 2023   | B     | 3.5 |
# | 2          | 101       | Fall 2023   | A     | 4.0 |
# | 2          | 102       | Spring 2023 | B     | 3.0 |
# | 3          | 103       | Fall 2023   | A     | 4.0 |
# | 3          | 104       | Spring 2023 | A     | 4.0 |
# | 3          | 106       | Spring 2023 | A     | 4.0 |
# | 3          | 108       | Fall 2023   | B     | 3.5 |
# | 4          | 103       | Fall 2023   | B     | 3.0 |
# | 4          | 104       | Spring 2023 | B     | 3.0 |
# +------------+-----------+-------------+-------+-----+
# Output:
# +------------+
# | student_id |
# +------------+
# | 1          |
# | 3          |
# +------------+
# Explanation:
# Alice (student_id 1) is a Computer Science major and has taken both Algorithms and Data Structures, receiving an A in both. She has also taken Machine Learning and Operating Systems as electives, receiving an A and B respectively.
# Bob (student_id 2) is a Computer Science major but did not receive an A in all required courses.
# Charlie (student_id 3) is a Mathematics major and has taken both Calculus and Linear Algebra, receiving an A in both. He has also taken Probability and Statistics as electives, receiving an A and B respectively.
# David (student_id 4) is a Mathematics major but did not receive an A in all required courses.
# Note: Output table is ordered by student_id in ascending order.

import pandas as pd

def find_top_scoring_students(students: pd.DataFrame, courses: pd.DataFrame, enrollments: pd.DataFrame) -> pd.DataFrame:
  # Get mandatory and elective courses per major
  mandatory = courses[courses['mandatory'].str.lower() == 'yes']
  elective = courses[courses['mandatory'].str.lower() == 'no']
  # Join students with their major's mandatory courses
  student_mandatory = students.merge(mandatory, on='major')[['student_id', 'course_id']]
  mandatory_count = student_mandatory.groupby('student_id')['course_id'].nunique().reset_index(name='mandatory_required')
  # Check mandatory: all A
  a_mandatory = enrollments[enrollments['grade'] == 'A'].merge(student_mandatory, on=['student_id', 'course_id'])
  a_mandatory_count = a_mandatory.groupby('student_id')['course_id'].nunique().reset_index(name='mandatory_done')
  # Check elective: grade A or B, at least 2
  student_elective = students.merge(elective, on='major')[['student_id', 'course_id']]
  ab_elective = enrollments[enrollments['grade'].isin(['A', 'B'])].merge(student_elective, on=['student_id', 'course_id'])
  ab_elective_count = ab_elective.groupby('student_id')['course_id'].nunique().reset_index(name='elective_done')
  # Average GPA >= 2.5 across all courses
  avg_gpa = enrollments.groupby('student_id')['GPA'].mean().reset_index(name='avg_gpa')
  # Combine
  result = students[['student_id']].merge(mandatory_count, on='student_id', how='left').fillna(0)
  result = result.merge(a_mandatory_count, on='student_id', how='left').fillna(0)
  result = result.merge(ab_elective_count, on='student_id', how='left').fillna(0)
  result = result.merge(avg_gpa, on='student_id', how='left').fillna(0)
  result = result[
    (result['mandatory_required'] == result['mandatory_done']) &
    (result['mandatory_required'] > 0) &
    (result['elective_done'] >= 2) &
    (result['avg_gpa'] >= 2.5)
  ][['student_id']]
  return result.sort_values('student_id').reset_index(drop=True)