На первый взгляд, это тривиальное преобразование после встраивания определения distances
:
predictions = []
for train_row in test_data:
distances = ...
predictions.append(train_tgt[distances])
становится
predictions = []
for train_row in test_data:
predictions.append(train_tgt[...])
, которое становится
predictions = [train_tgt[...] for train_row in test_data]