获取数据帧行中n个最大绝对值的索引

import heapq

def get_n_largest_indexes(data_row, n):
    # 1. 绝对值化数据行
    abs_data_row = [abs(x) for x in data_row]

    if n == 1:
        # 如果n等于1，直接找到绝对值最大的元素的索引
        max_index = abs_data_row.index(max(abs_data_row))
        return [max_index]

    else:
        # 使用堆排序获取绝对值最大的n个元素的索引
        max_indexes = []
        heap = []

        for i in range(len(abs_data_row)):
            if i < n:
                # 前n个元素直接加入堆中
                heapq.heappush(heap, (-abs_data_row[i], i))
            else:
                # 维护一个大小为n的最大堆
                heapq.heappushpop(heap, (-abs_data_row[i], i))

        while heap:
            max_indexes.append(heapq.heappop(heap)[1])

        return max_indexes

# 示例用法
data_row = [10, -5, 20, -15, 25, -30]
n = 3
result = get_n_largest_indexes(data_row, n)
print(result)