Visualization of monty hall problem

montyhall.py

from random import randrange
from matplotlib import use
use('AGG')
from matplotlib import pyplot as plt


class MontyHallproblem:

    def __init__(self, fname, tryN):
        self.fname = fname
        self.tryN = tryN

    def __enter__(self):
        return self

    def __change(self, player, m):
        for i in range(0, 3):
            if i != player and i != m:
                return i

    def __monty(self, prize, player):
        if prize == player:
            if player & 1:
                return randrange(0, 3, 2)
            else:
                return randrange(1, 3) if not player else randrange(0, 2)
        else:
            return self.__change(prize, player)

    def __game(self):
        prize = randrange(3)
        player = randrange(3)
        return player == prize, self.__change(player, self.__monty(
            prize, player)) == prize

    def config(self):
        ax = plt.figure().add_subplot(1, 1, 1)
        ax.set_title('Monty Hall problem')
        x = range(self.tryN)
        changed, not_changed = 0, 0
        changed_result, not_changed_result = [], []
        for _ in x:
            t1, t2 = self.__game()
            not_changed_result.append(not_changed)
            changed_result.append(changed)
            not_changed += t1
            changed += t2
        ax.set_xlabel('Number of trials')
        ax.set_ylabel('Number of times the prize door was selected')
        ax.plot(x, changed_result, label='Choose C')
        ax.plot(x, not_changed_result, label='Choose A')
        plt.legend(loc='lower right')

    def __exit__(self, exception_type, exception_value, traceback):
        plt.savefig(self.fname, antialiased=True)


if __name__ == '__main__':
    with MontyHallproblem('montyhalll.png', 1000) as m:
        m.config()