Projection Onto Probability Simplex

Problem

Projecting y onto the probability simplex (nonnegative entries summing to one) uses the sorted-threshold algorithm: sort y descending, find the largest rho with y_sorted[rho-1] + (1 - sum of the top rho) / rho > 0, then return max(y_i - theta, 0) with theta = (sum of the top rho - 1) / rho.

Given y, return the projected vector.

Hint

The threshold is the average excess of the active top block above one.

Test cases

case 1

in: [[0.5,0.5]]

expected: [0.5,0.5]

case 2

in: [[2,1,0]]

expected: [1,0,0]

case 3

in: [[-1,4,-2,1]]

expected: [0,1,0,0]

Starter codePython

def projection_onto_probability_simplex(y):
    # Your code here
    pass

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