如何对 Pandas 数据框进行二分搜索以获取列值的组合？

Question

Sorry if this is a simple question that the pandas documentation explains, but I've tried searching for how to do this and haven't had any luck.

I have a pandas datafame with several columns, and I want to be able to search for a particular row using binary search since my dataset is big and I'll be doing a lot of searches.

My data looks like this:

Name           Course   Week  Grade
-------------  -------  ----  -----
Homer Simpson  MATH001  1     97
Homer Simpson  MATH001  3     85
Homer Simpson  CSCI100  1     89
John McGuirk   MATH001  2     78
John McGuirk   CSCI100  1     100
John McGuirk   CSCI100  2     96

I want to be able to search my data quickly for a specific combination of name, course, and week. Each distinct combination of name, course, and week will have either zero or one row in the dataset. If there is a missing value for the combination of name, course, and week that I'm searching for, I want my search to return 0.

For instance, I would like to search for the value (John McGuirk, CSCI100, 1)

Is there a built in way to do this, or do I have to write my own binary search?

Update:

I tried doing this using the built-in way that was suggested by one of the commenters below, and I also tried doing it with a custom binary search that's written for my specific data, and another custom binary search that uses recursion to handle different columns than my specific example.

The data frame for these tests contains 10,000 rows. I put the timings below. Both binary searches performed better than using [...] to get rows. I'm far from a Python expert, so I'm not sure how well optimized my code is.

# Load data
from pandas import DataFrame, read_csv
import math
import pandas as pd
import time

file = 'grades.xlsx'
df = pd.read_excel(file)

# This was suggested by one of the commenters below
def get_grade(name, course, week):
    mask = (df.name.values == name) & (df.course.values == course) & (df.week.values == week)
    row = df[mask]
    if row.empty == False:
        return row.grade.values[0]
    else:
        return 0

# Binary search that is specific to my particular data
def get_grade_binary_search(name, course, week):
    lower = 0
    upper = len(df.index) - 1

    while lower <= upper:
        mid = math.floor((lower + upper) / 2)

        row_name = df.iat[mid, 0]            
        if name < row_name:
            upper = mid - 1
        elif name > row_name:
            lower = mid + 1
        else:
            row_course = df.iat[mid, 1]
            if course < row_course:
                upper = mid - 1
            elif course > row_course:
                lower = mid + 1
            else:
                row_week = df.iat[mid, 2]
                if week < row_week:
                    upper = mid - 1
                elif week > row_week:
                    lower = mid + 1
                else:
                    return df.iat[mid, 3]

    return 0    

# General purpose binary search
def get_grade_binary_search_recursive(search_value):
    lower = 0
    upper = len(df.index) - 1

    while lower <= upper:
        mid = math.floor((lower + upper) / 2)

        comparison = compare(search_value, 0, mid)

        if comparison < 0:
            upper = mid - 1
        elif comparison > 0:
            lower = mid + 1
        else:
            return df.iat[mid, len(search_value)]

# Utility method
def compare(search_value, search_column_index, df_value_index):      
    if search_column_index >= len(search_value):
        return 0

    if search_value[search_column_index] < df.iat[df_value_index, search_column_index]:
        return -1
    elif search_value[search_column_index] > df.iat[df_value_index, search_column_index]:
        return 1
    else:
        return compare(search_value, search_column_index + 1, df_value_index)

Here are the timings. I also printed the sum of the returned values from each search to verify that the same rows are getting returned.

# Non binary search
sum_of_grades = 0
start = time.time()   
for week in range(first_week, last_week + 1):
    for name in names:
        for course in courses:
            val = get_grade(name, course, week)
            sum_of_grades += val                
end = time.time()    
print('elapsed time: ', end - start)
print('sum of grades: ', sum_of_grades)

elapsed time: 26.130020141601562

sum of grades: 498724

# Binary search specific to this data
sum_of_grades = 0
start = time.time()    
for week in range(first_week, last_week + 1):
    for name in names:
        for course in courses:
            val = get_grade_binary_search(name, course, week)
            sum_of_grades += val

end = time.time()    
print('elapsed time: ', end - start)
print('sum of grades: ', sum_of_grades)

elapsed time: 4.4506165981292725

sum of grades: 498724

# Binary search with recursion
sum_of_grades = 0
start = time.time()
for week in range(first_week, last_week + 1):
    for name in names:
        for course in courses:
            val = get_grade_binary_search_recursive([name, course, week])
            sum_of_grades += val           
end = time.time()    
print('elapsed time: ', end - start)
print('sum_of_grades: ', sum_of_grades)

elapsed time: 7.559535264968872

sum_of_grades: 498724

Answer 1

Pandas has searchsorted .

From the Notes :

Binary search is used to find the required insertion points.