UCB Bandit Selection

Problem

Select an arm with the UCB1 rule.

Any arm with count 0 is selected first, smallest index winning. Otherwise pick the arm maximizing q_values[i] + c * sqrt(log(total) / counts[i]), where total is the sum of all counts. Ties go to the smallest index.

Hint

Unvisited arms always get pulled before any confidence bonus is computed.

Test cases

case 1

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

expected: 1

case 2

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

expected: 0

case 3

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

expected: 0

case 4

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

expected: 0

case 5

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

expected: 0

Starter codePython

def ucb_select(q_values, counts, c):
    # Your code here
    pass

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