archive/Google FooBar/doomsdayFuel.py

71 lines
1.7 KiB
Python

from fractions import Fraction
def div_row(l, row, m):
for i in range(len(l[row])): l[row][i] = l[row][i] / m
def add_row(l, row1, row2, m=1):
for i in range(len(l[row1])): l[row1][i] = l[row1][i] + l[row2][i] * m
def process(l):
# Converts matrix to fractions, and gets terminal states.
# Transforms L into (I - L^T) | (1,0,0,...)^T
new_l = [[0]*len(l) for x in range(len(l))]
terminal = []
for i in range(len(l)):
d = sum(l[i])
if d == 0: terminal.append(i)
for j in range(len(l)):
new_l[j][i] = -Fraction(l[i][j], 1 if d == 0 else d)
if i == j: new_l[i][j] = 1 + new_l[i][j]
for i in range(len(l)): new_l[i].append(Fraction(i == 0))
return (new_l, terminal)
def gauss_elim(l):
for i in range(len(l)): # Column
for j in range(len(l)): # Row
if i == j:
if l[j][i] == 1: continue
div_row(l, j, l[j][i])
else:
if l[j][i] == 0: continue
add_row(l, j, i, Fraction(-l[j][i], l[i][i]))
def gcd(a, b):
while b > 0:
a, b = b, a % b
return a
def lcm(a, b):
return a * b / gcd(a, b)
def simplify(p):
denom = [x.denominator for x in p]
curr_lcm = denom[0]
for i in range(1, len(denom)):
curr_lcm = lcm(curr_lcm, denom[i])
for i in range(len(p)):
p[i] = p[i].numerator * curr_lcm / denom[i]
p.append(curr_lcm)
def solution(l):
# L^n[i][j] = Probability of arriving at state j starting from state i in n steps
# Sum(L^n) from 0 to n = (I - L)^{-1}.
# Since we only want the first row, we can do some simplifications by
# row reducing with the transpose, and augment with the (1,0,0,0, ...)^T
l, terminal = process(l)
print(l)
gauss_elim(l)
probs = [x[-1] for i,x in enumerate(l) if i in terminal]
simplify(probs)
return probs