213 lines
6.0 KiB
Python
213 lines
6.0 KiB
Python
|
from fractions import Fraction
|
||
|
import unittest
|
||
|
|
||
|
def mAdd(left, right):
|
||
|
if len(left[0]) < len(right[0]):
|
||
|
short = left
|
||
|
else:
|
||
|
short = right
|
||
|
result = [[0 for i in range(len(short))] for j in range(len(short[0]))]
|
||
|
for i in range(len(result)):
|
||
|
for j in range(len(result[0])):
|
||
|
result[i][j] = left[i][j] + right[i][j]
|
||
|
return result
|
||
|
|
||
|
def mSub(left, right):
|
||
|
if len(left[0]) < len(right[0]):
|
||
|
short = left
|
||
|
else:
|
||
|
short = right
|
||
|
result = [[0 for i in range(len(short))] for j in range(len(short[0]))]
|
||
|
for i in range(len(result)):
|
||
|
for j in range(len(result[0])):
|
||
|
result[i][j] = left[i][j] - right[i][j]
|
||
|
return result
|
||
|
|
||
|
def mMul(left, right):
|
||
|
result = [[sum(a * b for a, b in zip(left_row, right_col))
|
||
|
for right_col in zip(*right)]
|
||
|
for left_row in left]
|
||
|
return result
|
||
|
|
||
|
def makeIdentityM(s):
|
||
|
result = [[0 for i in range(s)] for j in range(s)]
|
||
|
for i, row in enumerate(result):
|
||
|
for j, col in enumerate(row):
|
||
|
if i == j:
|
||
|
result[i][j] = 1
|
||
|
return result
|
||
|
|
||
|
def makeRegAndSTD(m):
|
||
|
shift = 0
|
||
|
# Point Terminal states to themselves and move to top of 2d array
|
||
|
for i, row in enumerate(m):
|
||
|
void = True
|
||
|
for col in row:
|
||
|
if col != 0:
|
||
|
void = False
|
||
|
if void == True:
|
||
|
row[i] = 1
|
||
|
m.insert(shift, row)
|
||
|
m.pop(i+1)
|
||
|
shift += 1
|
||
|
# Shift all elements relative to the nmber of terminal states found
|
||
|
for i, row in enumerate(m):
|
||
|
m[i] = m[i][-shift:] + m[i][:-shift]
|
||
|
|
||
|
# Extract sections of the matrix for use in calculating F
|
||
|
top = m[:shift]
|
||
|
bottom = m[shift:]
|
||
|
I = [[0 for i in range(len(m))] for j in range(shift)]
|
||
|
R = [[0 for i in range(len(m))] for j in range(len(m)-shift)]
|
||
|
Q = [[0 for i in range(len(m))] for j in range(len(m)-shift)]
|
||
|
for i, col in enumerate(top):
|
||
|
I[i] = col[:shift]
|
||
|
for i, col in enumerate(bottom):
|
||
|
R[i] = col[:shift]
|
||
|
Q[i] = col[shift:]
|
||
|
return [m, I, R, Q]
|
||
|
|
||
|
def invertMatrix(AM, IM):
|
||
|
for fd in range(len(AM)):
|
||
|
fdScaler = 1.0 / AM[fd][fd]
|
||
|
for j in range(len(AM)):
|
||
|
AM[fd][j] *= fdScaler
|
||
|
IM[fd][j] *= fdScaler
|
||
|
for i in list(range(len(AM)))[0:fd] + list(range(len(AM)))[fd+1:]:
|
||
|
crScaler = AM[i][fd]
|
||
|
for j in range(len(AM)):
|
||
|
AM[i][j] = AM[i][j] - crScaler * AM[fd][j]
|
||
|
IM[i][j] = IM[i][j] - crScaler * IM[fd][j]
|
||
|
return IM
|
||
|
|
||
|
def gcd(x, y):
|
||
|
while y:
|
||
|
x, y = y, x % y
|
||
|
return x
|
||
|
|
||
|
def transformProbabilities(m):
|
||
|
result = [0 for i in range(len(m))]
|
||
|
numerators = [0 for i in range(len(m))]
|
||
|
denomenators = [0 for i in range(len(m))]
|
||
|
|
||
|
for i, prob in enumerate(m):
|
||
|
x = Fraction(prob).limit_denominator(1000)
|
||
|
if str(x).find('/') != -1:
|
||
|
y = str(x).split('/')
|
||
|
numerators[i] = int(y[0])
|
||
|
denomenators[i] = int(y[1])
|
||
|
else:
|
||
|
numerators[i] = int(x)
|
||
|
|
||
|
lcm = denomenators[0]
|
||
|
# Protect by divide 0
|
||
|
if lcm == 0:
|
||
|
lcm = 1
|
||
|
for i in denomenators[1:]:
|
||
|
lcm = lcm*i/gcd(lcm, i)
|
||
|
# Make sure lcm is not 0
|
||
|
if lcm == 0:
|
||
|
lcm = 1
|
||
|
|
||
|
for i, num in enumerate(numerators):
|
||
|
if denomenators[i] != 0 and denomenators[i] != lcm:
|
||
|
result[i] = num * (lcm / denomenators[i])
|
||
|
else:
|
||
|
result[i] = num
|
||
|
|
||
|
result.append(lcm)
|
||
|
return result
|
||
|
|
||
|
def solution(m):
|
||
|
# Check if our matrix is 1 x 1
|
||
|
try:
|
||
|
m[0][1]
|
||
|
except IndexError:
|
||
|
return [1,1]
|
||
|
|
||
|
# Check if S0 is Terminal State
|
||
|
isTerminal = True
|
||
|
for item in m[0][1:]:
|
||
|
if item != 0:
|
||
|
isTerminal = False
|
||
|
if isTerminal == True:
|
||
|
result = [1]
|
||
|
for i, row in enumerate(m[1:]):
|
||
|
void = True
|
||
|
for col in row:
|
||
|
if col != 0:
|
||
|
void = False
|
||
|
if void == True:
|
||
|
result.append(0)
|
||
|
result.append(1)
|
||
|
return result
|
||
|
|
||
|
# Write fractions to matrix (for help with math)
|
||
|
for i, row in enumerate(m):
|
||
|
den = 0
|
||
|
for j, col in enumerate(row):
|
||
|
den += m[i][j]
|
||
|
for j, col in enumerate(row):
|
||
|
if m[i][j] != 0 and den != 0:
|
||
|
m[i][j] = Fraction(m[i][j], den)
|
||
|
|
||
|
# Returns [m, I, R, Q]
|
||
|
helpers = makeRegAndSTD(m)
|
||
|
iq = mSub(helpers[1], helpers[3])
|
||
|
F = invertMatrix(iq, makeIdentityM(len(iq)))
|
||
|
FR = mMul(F, helpers[2])
|
||
|
return transformProbabilities(FR[0])
|
||
|
|
||
|
|
||
|
class TestDoomsdayFuel(unittest.TestCase):
|
||
|
def test1(self):
|
||
|
test_input = [
|
||
|
[0, 2, 1, 0, 0],
|
||
|
[0, 0, 0, 3, 4],
|
||
|
[0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0]
|
||
|
]
|
||
|
self.assertEqual(solution(test_input),
|
||
|
[7, 6, 8, 21])
|
||
|
|
||
|
def test2(self):
|
||
|
test_input = [
|
||
|
[0, 1, 0, 0, 0, 1],
|
||
|
[4, 0, 0, 3, 2, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0]
|
||
|
]
|
||
|
self.assertEqual(solution(test_input), [0, 3, 2, 9, 14])
|
||
|
|
||
|
def test3(self):
|
||
|
test_input = [
|
||
|
[0, 1, 0, 0, 0, 1],
|
||
|
[1, 0, 0, 1, 1, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0],
|
||
|
[0, 0, 0, 0, 0, 0]
|
||
|
]
|
||
|
self.assertEqual(solution(test_input), [0, 1, 1, 3, 5])
|
||
|
|
||
|
def test4(self):
|
||
|
test_input = [
|
||
|
[1, 1, 0, 1],
|
||
|
[1, 1, 0, 0],
|
||
|
[0, 0, 0, 0],
|
||
|
[0, 0, 0, 0],
|
||
|
]
|
||
|
self.assertEqual(solution(test_input), [0, 1, 1])
|
||
|
|
||
|
def test5(self):
|
||
|
test_input = [
|
||
|
[0, 1, 0],
|
||
|
[1, 1, 1],
|
||
|
[0, 0, 0]
|
||
|
]
|
||
|
self.assertEqual(solution(test_input), [1, 0, 1])
|
||
|
|
||
|
unittest.main()
|