Return to Answer

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

class Coin():
    
    def __init__(self):
        self.a = 1
        self.b = 1
    def draw(self):
        return beta(self.a, self.b).rvs(1)
    def update(self, flip):
        if flip>0:
            self.a+=1
        else:
            self.b+=1


#Unknown to us
np.random.seed(19920908)
coin1 = binom(p=0.4, n=1)
coin2 = binom(p=0.6, n=1)


model1 = Coin()
model2 = Coin()

for i in range(100):

    draw1 = model1.draw()
    draw2 = model2.draw()

    if draw1>draw2:
        flip = coin1.rvs()
        model1.update(flip)
    else:
        flip = coin2.rvs()
        model2.update(flip)


        
x = np.linspace(0,1,101)

plt.plot(x, beta(model1.a, model1.b).pdf(x))
plt.plot(x, beta(model2.a, model2.b).pdf(x))

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

class Coin():
    
    def __init__(self):
        self.a = 1
        self.b = 1
    def draw(self):
        return beta(self.a, self.b).rvs(1)
    def update(self, flip):
        if flip>0:
            self.a+=1
        else:
            self.b+=1
    def __str__(self):
        return f"{self.a}:{self.b}={self.a/(self.a+self.b):.3f}"



#Unknown to us
np.random.seed(19920908)
coin1 = binom(p=0.4, n=1)
coin2 = binom(p=0.6, n=1)


model1 = Coin()
model2 = Coin()

for i in range(100):

    draw1 = model1.draw()
    draw2 = model2.draw()

    if draw1>draw2:
        flip = coin1.rvs()
        model1.update(flip)
    else:
        flip = coin2.rvs()
        model2.update(flip)


        
x = np.linspace(0,1,101)

plt.plot(x, beta(model1.a, model1.b).pdf(x))
plt.plot(x, beta(model2.a, model2.b).pdf(x))
print(model1,model2)

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

%matplotlib inline

coin1 = binom(p = 0.4, n=1)
coin2 = binom(p = 0.6, n=1)


Aa = 1
Ab = 1
Ba = 1
Bb = 1
n_flips = 100

A = beta(Aa, Ab)
B = beta(Ba, Bb)

for i in range(100):
    
     draw_A = A.rvs()
    draw_B = B.rvs()
    
    if draw_A>draw_B:
        flip = coin1.rvs()
        if flip>0:
            Aa+=1
            A = beta(Aa, Ab)
        else:
            Ba+=1
            A = beta(Aa, Ab)
    else:
        flip = coin2.rvs()
        if flip>0:
            Ba+=1
            B = beta(Ba, Bb)
        else:
            Bb+=1
            B = beta(Ba, Bb)
        

 
x = np.linspace(0,1, 101)

plt.plot(x, A.pdf(x))
plt.plot(x,B.pdf(x))

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

import numpy as np
from scipy.stats import beta, binom
import matplotlib.pyplot as plt

class Coin():
    
    def __init__(self):
        self.a = 1
        self.b = 1
    def draw(self):
        return beta(self.a, self.b).rvs(1)
    def update(self, flip):
        if flip>0:
            self.a+=1
        else:
            self.b+=1


#Unknown to us
np.random.seed(19920908)
coin1 = binom(p=0.4, n=1)
coin2 = binom(p=0.6, n=1) 


model1 = Coin()
model2 = Coin()

for i in range(100):

    draw1 = model1.draw()
    draw2 = model2.draw()

    if draw1>draw2:
        flip = coin1.rvs()
        model1.update(flip)
    else:
        flip = coin2.rvs()
        model2.update(flip) 


        
x = np.linspace(0,1,101)

plt.plot(x, beta(model1.a, model1.b).pdf(x))
plt.plot(x, beta(model2.a, model2.b).pdf(x))

EDIT: I'm leaving this up for posterity, but this probably should not be the accepted answer. In the comments I note that assuming the coins have a bias of 0.5 and 0.51, then OP's strategy selects the correct coin 58% of the time. The bandit is (ironically) a coin flip.

Nothing frustrates me more than when someone tells you to do the "optimal" thing without telling you the criteria over which to optimize. That being said, I'm betting that since it was an interview, they intended for you to determine what you wanted to optimize for.

Nothing frustrates me more than when someone tells you to do the "optimal" thing without telling you the criteria over which to optimize. That being said, I'm betting that since it was an interview, they intended for you to determine what you wanted to optimize for.