Probability calculation for civ3 battle round

civ3.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-


"""
A problem from civ3 game. Let's say attacker has m hp, defender has n hp. In
each round, attacker will inflict damage of 1hp with probability of p, or lose
1hp with probability of 1-p. We ask the probability of final outcome.
"""


import sys
import numpy as np
import scipy.sparse as sp


def main(args):
    p, m, n = args
    p, m, n = float(p), int(m), int(n)
    print "prob={0} attacker-hp={1} defender-hp={2}".format(p, m, n)

    conf = [(i, j) for i in xrange(m + 1) for j in xrange(n + 1)]
    N = len(conf)
    I = dict((x, i) for i, x in enumerate(conf))   # configuration i
    P = sp.lil_matrix((N, N), dtype=float)  # transition matrix
    E = np.zeros(N, dtype=float)  # pi
    for x, y in conf:  # build transition matrix
        i = I[x, y]
        if x and y != 0:
            j = I[x - 1, y]
            P[j, i] = 1 - p
        if y and x != 0:
            j = I[x, y - 1]
            P[j, i] = p
        if not x or not y:
            P[i, i] = 1

    E[I[m, n]] = 1
    for i in xrange(m + n - 1):
        print "Round #{0}:".format(i)
        E = P.dot(E)
        print E

    print "*" * 20
    attacker_dies = defender_dies = 0
    for i, v in enumerate(E):
        if v > 1e-10:
            print conf[i], "prob={0}".format(v)
            c = conf[i]
            if c[0] == 0:
                attacker_dies += v
            if c[-1] == 0:
                defender_dies += v
    print "ATTACKER DIES WITH PROB={0}".format(attacker_dies)
    print "DEFENDER DIES WITH PROB={0}".format(defender_dies)


if __name__ == '__main__':
    main(sys.argv[1:])