final function

separating_process.py

# Parse combineed grades
def parse_grades(grades_string):
    GRADES = ['PK', 'K', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', 'Ungraded']

    # Remove & for grades list
    grades_string = grades_string.replace(' &', ',')

    # Grades list - will add to separated grade string to grades
    grades = []

    # split strings based on ','
    string_list = grades_string.split(',')

    # look for sections of list with '-'
    dash = "-"
    for i in range(len(string_list)):
        clean_string = string_list[i].strip()
        if dash in clean_string:
            #  split using '-', loop and add to grades variable
            start_grade, end_grade = clean_string.split(dash)
            grades += GRADES[GRADES.index(start_grade) : GRADES.index(end_grade)+ 1]
        else:
            # add string to grades
            grades.append(clean_string)
    return grades
 
# Create dictionary  
def test_complete_dataset(unique_grades_combination):
    # create a loop that goes thru dataset and invoke parse_grades with each element
    separated_grades_list = []
    for i in unique_grades_combination:
        separated_grades_list.append(parse_grades(i))

    dictionary_grade_list = dict(zip(unique_grades_combination, separated_grades_list))

    return dictionary_grade_list

# Create new columns with separated grade
df['Clean_Grades'] = df['Grades'].map(dictionary)

# Schools category List
high_school = ['9', '10', '11', '12']
middle_school = ['6', '7', '8']
elementary = ['K', '1', '2', '3', '4', '5']
pre_k = ['PK']

# Mapping each category
set1 = set(high_school)
df['High School (9-12)'] = df['Clean_Grades'].apply(lambda x: any([k in x for k in set1]))

set2 = set(middle_school)
df['Middle School (6-8)'] = df['Clean_Grades'].apply(lambda x: any([k in x for k in set2]))

set3 = set(elementary)
df['Elementary (K-5)'] = df['Clean_Grades'].apply(lambda x: any([k in x for k in set3]))

set4 = set(pre_k)
df['Pre-Kindergarten (PK)'] = df['Clean_Grades'].apply(lambda x: any([k in x for k in set4]))