如何获得topk的值及其指数(2D)？

def get_topk_values_and_indices_2d(data, k):
    # 对数据集按照需要进行排序
    sorted_data = sorted(data, key=lambda x: x[0], reverse=True)  # 以第一列的值进行降序排序

    # 获取topk的值及其指数
    topk_values = [item[0] for item in sorted_data[:k]]
    topk_indices = [item[1] for item in sorted_data[:k]]

    return topk_values, topk_indices

# 示例数据集
data = [[5, 0], [3, 1], [8, 2], [2, 3], [9, 4], [1, 5]]
k = 3

# 调用函数获取topk的值及其指数
topk_values, topk_indices = get_topk_values_and_indices_2d(data, k)

# 打印结果
print("Topk values:", topk_values)
print("Topk indices:", topk_indices)