TFT Variable Selection Weights

Problem

Compute TFT-style variable selection weights: apply a softmax to the scores and return the top_k weights in descending order.

Return an empty list for empty scores or top_k <= 0. The exponential shift by the maximum keeps the computation stable.

Hint

Softmax weights always sum to 1 before truncation.

case 1

in: [[1,2,3],2]

expected: [0.6652409557748218,0.24472847105479764]

case 2

in: [[0,0],1]

expected: [0.5]

case 3

in: [[],2]

expected: []

case 4

in: [[5,5,5],3]

expected: [0.3333333333333333,0.3333333333333333,0.3333333333333333]

Starter codePython

import math
def tft_variable_selection_weights(scores, top_k):
    # Your code here
    pass

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