sts_markov.py

#!/usr/bin/env python

class QMN:
  def __init__(self, streak, total, chance):
    self.streak = streak
    self.total = total
    self.chance = chance

  def childNodes(self):
    enemyChance = 0.1 * (self.streak + 1)
    yield QMN(self.streak + 1, self.total, (1 - enemyChance) * self.chance)
    yield QMN(0, self.total + 1, enemyChance * self.chance)

def frequencyCount(items):
  freq = dict()
  for item in items:
    if item[0] in freq:
      freq[item[0]] = freq[item[0]] + item[1]
    else:
      freq[item[0]] = item[1]
  return freq

def processMarkov(nodes, agreggationFn, maxDepth):
  if maxDepth <= 0:
    return frequencyCount((agreggationFn(x), x.chance) for x in nodes)
  nextNodes = []
  for node in nodes:
    nextNodes.extend(node.childNodes())
  return processMarkov(nextNodes, agreggationFn, maxDepth - 1)

def cumulativeProbabilities(prob):
  acc = 0
  for i in xrange(len(prob.keys())):
    acc += prob[i]
    yield "{0:.0%}".format(acc)

if __name__ == '__main__':
  for i in range(1, 7):
    print "Probabilities for %d QMNs" % i
    result = processMarkov([QMN(0, 0, 1)], lambda x: x.total, i)
    print list(cumulativeProbabilities(result))
    print "\n"