Python二维数组按指定列排序异常：IMDb数据集startYear排序失效

我正在开发Python函数，使用多种排序算法对IMDb数据集按测试用例指定的列排序（如测试用例1需按startYear列排序），但运行首个测试后发现排序算法未按startYear排序，反而按movieId（tconst）排序，导致测试失败。以下是相关代码及测试结果：

核心排序函数 #

def sorting_algorithms(file_path, columns, select):
    df = pd.read_csv(file_path)
    column_vals = [0] + [df.columns.get_loc(col) for col in columns]
    # Convert the dataframe to a 2D array of values
    data = df.iloc[:, column_vals].values.tolist()
    
    if(select==1):
        start_time = time.time()
        output_list = insertion_sort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==2):
        start_time = time.time()
        output_list = selection_sort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==3):
        start_time = time.time()
        output_list = quicksort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==4):
        start_time = time.time()
        output_list = heap_sort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==5):
        start_time = time.time()
        output_list = shell_sort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==6):
        start_time = time.time()
        output_list = merge_sort(data, column_vals)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]

selection_sort排序算法 #

def selection_sort(arr, columns):
    for i in range(len(arr)):
        min_idx = i
        for j in range(i+1, len(arr)):
            for col in columns:
                if arr[j][col] < arr[min_idx][col]:
                    min_idx = j
                    break
                elif arr[j][col] > arr[min_idx][col]:
                    break
        arr[i], arr[min_idx] = arr[min_idx], arr[i]
    return arr

测试函数 #

def testcase_1_1():
    global total
    global f
    sorted_testcase_1_1 = sorting_algorithms("testcases_1_2_df.csv", ['startYear'], 2)
    print("Sorted list of movie IDs for test case 1_1: ", sorted_testcase_1_1[1])

    df = pd.read_csv('imdb_dataset.csv', low_memory=False, delimiter=',')
    df = df[df['tconst'].isin(sorted_testcase_1_1[1])].set_index('tconst').loc[sorted_testcase_1_1[1]].reset_index()

    print("Sorted 2D array based on startYear column:\n", df[['tconst', 'startYear']].values)

    prev_start_year = None
    for index, row in df.iterrows():
        start_year = int(row['startYear'])
        print("prev_start_year: {}, start_year: {}".format(prev_start_year, start_year))
        if prev_start_year and start_year < prev_start_year:
            print("TestCase 1_1 failed")
            f += 1
            return 0
        prev_start_year = start_year

    print("TestCase 1_1 Passed and your Algorithm Time Complexity = {}".format(sorted_testcase_1_1[0]))
    total += 1
    return sorted_testcase_1_1

测试失败输出 #

prev_start_year: None, start_year: 2020
prev_start_year: 2020, start_year: 2015
TestCase 1_1 failed

问题原因 #

核心bug出在column_vals的构造逻辑：

column_vals = [0] + [df.columns.get_loc(col) for col in columns]

你把**第0列（tconst列）**加到了排序依据的列列表最前面，导致排序算法优先比较tconst列。由于数据中tconst是唯一值，排序逻辑完全变成按tconst排序，完全忽略了指定的startYear列。

修复方案 #

1. 修改核心排序函数 #

将排序用的列索引和返回用的tconst列分离，确保排序时只使用指定列：

def sorting_algorithms(file_path, columns, select):
    df = pd.read_csv(file_path)
    # 获取指定排序列在原DataFrame中的索引
    sort_col_indices = [df.columns.get_loc(col) for col in columns]
    # 提取包含tconst（第0列）和排序列的数据
    data = df.iloc[:, [0] + sort_col_indices].values.tolist()
    # 构造排序用的相对索引：数据列表中排序列的位置是1、2...（第0位是tconst）
    sort_rel_indices = list(range(1, len(sort_col_indices)+1))
    
    if(select==1):
        start_time = time.time()
        output_list = insertion_sort(data, sort_rel_indices)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    if(select==2):
        start_time = time.time()
        output_list = selection_sort(data, sort_rel_indices)
        end_time = time.time()
        time_in_seconds = end_time - start_time
        return [time_in_seconds, list(map(lambda x: x[0], output_list))]
    # 其他排序算法分支同理替换传入的sort_rel_indices

2. 验证排序逻辑 #

修改后，selection_sort会优先比较数据列表中的第1列（对应原DataFrame的startYear），只有当startYear相等时才会继续比较后续指定列，完全符合测试要求。

可选：用Pandas内置排序做对照 #

如果需要快速验证测试逻辑是否正确，可以先用Pandas内置排序实现一个基准版本：

def sorting_algorithms(file_path, columns, select):
    df = pd.read_csv(file_path)
    start_time = time.time()
    sorted_df = df.sort_values(by=columns).reset_index(drop=True)
    end_time = time.time()
    return [end_time - start_time, sorted_df['tconst'].tolist()]

内容的提问来源于stack exchange，提问作者Yacine