exact.py

"""Small standard-library exact affine/rational linear-algebra helper."""
from fractions import Fraction as F

class Expr:
    def __init__(self,value=0):
        self.d=dict(value.d) if isinstance(value,Expr) else ({value:F(1)} if isinstance(value,str) else {'':F(value)})
        self.d={k:v for k,v in self.d.items() if v}
    def __add__(self,x):
        x=Expr(x);y=Expr(self)
        for k,v in x.d.items():y.d[k]=y.d.get(k,F(0))+v
        y.d={k:v for k,v in y.d.items() if v};return y
    __radd__=__add__
    def __neg__(self):return self*-1
    def __sub__(self,x):return self+-Expr(x)
    def __rsub__(self,x):return Expr(x)+-self
    def __mul__(self,x):
        if isinstance(x,Expr):
            assert not any(k for k in x.d);x=x.d.get('',0)
        y=Expr();y.d={k:v*F(x) for k,v in self.d.items() if v*F(x)};return y
    __rmul__=__mul__
    def __truediv__(self,x):return self*F(1,x)
    def __eq__(self,x):return self.d==Expr(x).d
    def __hash__(self):return hash(tuple(sorted(self.d.items())))
    def coeff(self,x):return self.d.get(next(iter(x.d)),F(0))
    def subs(self,m):
        y=Expr(self.d.get('',0))
        for k,v in self.d.items():
            if k:y+=Expr(m.get(Expr(k),Expr(k)))*v
        return y
    @property
    def is_Integer(self):return not any(k for k in self.d) and self.d.get('',F(0)).denominator==1
    def __int__(self):assert self.is_Integer;return int(self.d.get('',0))
    def __str__(self):
        out=[]
        for k,v in self.d.items():
            if not k:out.append(str(v))
            elif v==1:out.append(k)
            elif v==-1:out.append('-'+k)
            else:out.append(str(v)+'*'+k)
        return ' + '.join(out).replace('+ -','- ') or '0'
    __repr__=__str__

def symbols(s):
    if ':' in s:
        pre,end=s.split(':');i=0
        while i<len(pre) and not pre[i].isdigit():i+=1
        return tuple(Expr(pre[:i]+str(n)) for n in range(int(pre[i:]),int(end)))
    vals=tuple(Expr(x) for x in s.split());return vals[0] if len(vals)==1 else vals
Rational=F
Basic=Expr
def expand(x):return Expr(x)
factor=expand
simplify=expand
def diff(x,s):return Expr(x).coeff(s)

class Matrix:
    def __init__(self,rows):
        rows=list(rows)
        if rows and not isinstance(rows[0],(list,tuple,Matrix)):rows=[[x] for x in rows]
        self.a=[[F(x) for x in (r.a[0] if isinstance(r,Matrix) else r)] for r in rows]
        self.n=len(self.a);self.m=len(self.a[0]) if self.n else 0
    @property
    def shape(self):return self.n,self.m
    @property
    def T(self):return Matrix(list(zip(*self.a)))
    def tolist(self):return self.a
    def __getitem__(self,key):
        r,c=key
        if isinstance(r,int) and isinstance(c,int):return self.a[r][c]
        rr=[r] if isinstance(r,int) else (list(range(self.n))[r] if isinstance(r,slice) else r)
        cc=[c] if isinstance(c,int) else (list(range(self.m))[c] if isinstance(c,slice) else c)
        return Matrix([[self.a[i][j] for j in cc] for i in rr])
    def rref(self):
        a=[r[:] for r in self.a];piv=[];r=0
        for c in range(self.m):
            k=next((i for i in range(r,self.n) if a[i][c]),None)
            if k is None:continue
            a[r],a[k]=a[k],a[r];v=a[r][c];a[r]=[x/v for x in a[r]]
            for i in range(self.n):
                if i!=r:
                    v=a[i][c];a[i]=[x-v*y for x,y in zip(a[i],a[r])]
            piv.append(c);r+=1
            if r==self.n:break
        return Matrix(a),tuple(piv)
    def rank(self):return len(self.rref()[1])
    def nullspace(self):
        r,p=self.rref();out=[]
        for f in range(self.m):
            if f not in p:
                v=[F(0)]*self.m;v[f]=1
                for i,c in enumerate(p):v[c]=-r.a[i][f]
                out.append(Matrix(v))
        return out
    def col_join(self,b):assert self.m==b.m;return Matrix(self.a+b.a)
    @staticmethod
    def hstack(*ms):return Matrix([[x for m in ms for x in m.a[i]] for i in range(ms[0].n)])
    def __mul__(self,b):
        assert self.m==b.n
        return Matrix([[sum(self.a[i][k]*b.a[k][j] for k in range(self.m)) for j in range(b.m)] for i in range(self.n)])
    def det(self):
        assert self.n==self.m;a=[r[:] for r in self.a];d=F(1)
        for c in range(self.m):
            k=next((i for i in range(c,self.n) if a[i][c]),None)
            if k is None:return F(0)
            if k!=c:a[c],a[k]=a[k],a[c];d=-d
            v=a[c][c];d*=v
            for i in range(c+1,self.n):
                q=a[i][c]/v;a[i]=[x-q*y for x,y in zip(a[i],a[c])]
        return d
MatrixBase=Matrix
def linsolve(pair):
    a,b=pair;r,p=Matrix.hstack(a,b).rref();assert len(p)==a.m and a.m not in p
    return {tuple(r.a[i][-1] for i in range(a.m))}