蚂蚁春招0316code最后一题

import heapq
input_A = input().split()
input_B = input().split()
input_C = input().split()
n, k = int(input_A[0]), int(input_A[1])
index_value_pairs = [(i, int(input_B[i])) for i in range(n)]
index_value_pairs.sort(key=lambda x: x[1])
values = [int(i) for i in input_C]
result_dict = {i: [values[i]] for i in range(n)}
for i in range(n):
    next_index = i + 1
    while next_index < n and -k <= index_value_pairs[i][1] - index_value_pairs[next_index][1]:
        current_index = index_value_pairs[i][0]
        next_element_index = index_value_pairs[next_index][0]
        heapq.heappush(result_dict[current_index], values[next_element_index])
        heapq.heappush(result_dict[next_element_index], values[current_index])
        next_index += 1
total_sum = sum(result_dict[i][0] for i in range(n))
print(total_sum)