python - 获取列表中 N 个最大值的索引，使这些值符合给定条件

Question

Ids = [ i for i, v in heapq.nlargest(RN, enumerate(score_real_test), key=operator.itemgetter(1))] 

这将给出列表 score_real_test 中 RN 最大值的索引。是否可以获取“在列表 score_real_test 中并满足布尔条件 COND”的 RN 最大值的索引？

Ids 应该包含 RN 索引。

编辑：现在我使用这个解决方案，但它不是最好的：

score_real_test_2 = np.sort( [ v for i,v in enumerate(score_real_test) if pred_real_test[i] == NOVEL ] )
score_real_test_2 = score_real_test_2[len(score_real_test_2)-RN:]
large_dist_ids = [i for i in range(len(score_real_test)) if score_real_test[i] in score_real_test_2]

Answer 1

If you want to keep this is a one-liner, you can simply replace enumerate by a list comprehension(or generator expression) that pre-filters your data based on the condition.

Like so:

Ids = [i for i, v in heapq.nlargest(RN, [(j,v) in enumerate(score_real_test) if pred_real_test[i] == NOVEL], key=operator.itemgetter(1))] 

Or if you want to keep it clearer, you can add a filtering step before the selection(still no need for pre-sorting):

tmp = [item for item in enumerate(score_real_test) if pred_real_test[item[0]] == NOVEL]
Ids = [i for i, v in heapq.nlargest(RN, tmp, key=operator.itemgetter(1))]