HardNLPnlp-275

Label-Balanced Example Selection

Problem

Select one few-shot example per label. examples is a list of [id, label, vector]; for each distinct label, in first-seen order, choose the example with the highest cosine similarity to query, keeping the earliest on ties.

Return the chosen ids in that label order.

Hint

Track first-seen label order and store the best (similarity, index, id) per label.

Test cases

case 1

in: [[1,0],[["e1","a",[1,0]],["e2","b",[0,1]],["e3","a",[0.5,0]],["e4","b",[2,0]]]]

expected: ["e1","e4"]

case 2

in: [[1,0],[]]

expected: []

case 3

in: [[0,0],[["e1","a",[1,0]],["e2","b",[0,1]]]]

expected: ["e1","e2"]

case 4

in: [[1,0],[["x","a",[0,1]],["y","a",[1,0]]]]

expected: ["y"]

case 5

in: [[1,0],[["p","z",[0,1]],["q","a",[1,0]]]]

expected: ["p","q"]

Starter codePython

import math
def label_balanced_selection(query, examples):
    # Your code here
    pass

5 test cases · run your code against them in the editor