Bandit Sample Average Update

Problem

Update an arm's sample-average estimate after observing a reward.

Increment counts[a], then set Q[a] += (reward - Q[a]) / counts[a]. Return [new_Q, new_counts] without mutating the inputs.

Hint

The incremental formula keeps the exact running mean without storing every reward.

case 1

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

expected: [[2,0],[1,0]]

case 2

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

expected: [[3,0],[2,0]]

case 3

in: [[1,3,5],[2,1,4],2,0]

expected: [[1,3,4],[2,1,5]]

Starter codePython

def sample_average_update(Q, counts, a, reward):
    # Your code here
    pass

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