helpin' andy

itersolve.py

import numpy as np
import random
import matplotlib.pyplot as plt

def dirichletPop(dim):
    transMat = np.random.dirichlet(np.ones(dim),size=dim)
    return transMat

def n0Pop(dim):
    n0 = (1/dim) * np.ones((1,dim))
    return n0

def getChoice(dim):
    choice = random.randint(0,dim-1)
    return choice

def solve(transMat, nX, dim):
    choiceList = []
    while True:
        choice = getChoice(dim)
        choiceList.append(choice)
        nX[0,choice] = 0
        nX = np.dot(nX, transMat)
        if np.sum(nX) == 0:
            return choiceList

def iterSolve(iterations, transMat, dim):
    choiceCtList = []
    for i in range(iterations):
        nXnew = n0Pop(dim)
        choiceList = solve(transMat, nXnew, dim)
        choiceCtList.append(len(choiceList))
    return choiceCtList

def main():
    iterations = 100
    dim = 5
    transMat = dirichletPop(dim)
    choiceCtList = iterSolve(iterations, transMat, dim)
    print(choiceCtList)
    
main()