algebra.py

from fractions import Fraction as F

class Matrix:
    def __init__(self, rows):
        if isinstance(rows, Matrix): rows=rows.tolist()
        rows=list(rows)
        if rows and not isinstance(rows[0], (list,tuple)): rows=[[v] for v in rows]
        self.a=[[F(v) for v in r] for r in rows]
        self.rows=len(self.a);self.cols=len(self.a[0]) if self.a else 0
        assert all(len(r)==self.cols for r in self.a)
    def tolist(self):return [r.copy() for r in self.a]
    def __iter__(self):return iter([v for r in self.a for v in r])
    def __getitem__(self,key):
        if isinstance(key,tuple):
            i,j=key
            if isinstance(j,slice):return Matrix([self.a[i][j]])
            return self.a[i][j]
        return list(self)[key]
    @property
    def T(self):return Matrix(list(map(list,zip(*self.a))))
    def col_join(self,B):
        assert self.cols==B.cols
        return Matrix(self.a+B.a)
    def __mul__(self,B):
        assert self.cols==B.rows
        return Matrix([[sum(self.a[i][k]*B.a[k][j] for k in range(self.cols)) for j in range(B.cols)] for i in range(self.rows)])
    def rref(self):
        a=self.tolist();pivots=[];row=0
        for col in range(self.cols):
            candidates=[i for i in range(row,self.rows) if a[i][col]]
            if not candidates:continue
            p=candidates[0];a[row],a[p]=a[p],a[row];v=a[row][col]
            a[row]=[x/v for x in a[row]]
            for i in range(self.rows):
                if i!=row:
                    v=a[i][col];a[i]=[x-v*y for x,y in zip(a[i],a[row])]
            pivots.append(col);row+=1
            if row==self.rows:break
        return Matrix(a),pivots
    def rank(self):return len(self.rref()[1])
    def nullspace(self):
        rr,pivots=self.rref();free=[j for j in range(self.cols) if j not in pivots];out=[]
        for f in free:
            v=[F(0)]*self.cols;v[f]=F(1)
            for row,p in enumerate(pivots):v[p]=-rr.a[row][f]
            out.append(Matrix(v))
        return out
    def det(self):
        assert self.rows==self.cols
        a=self.tolist();d=F(1)
        for c in range(self.cols):
            ps=[i for i in range(c,self.rows) if a[i][c]]
            if not ps:return F(0)
            p=ps[0]
            if p!=c:a[c],a[p]=a[p],a[c];d=-d
            pivot=a[c][c];d*=pivot
            for i in range(c+1,self.rows):
                f=a[i][c]/pivot
                for j in range(c,self.cols):a[i][j]-=f*a[c][j]
        return d
    def gauss_jordan_solve(self,B):
        assert self.rows==B.rows
        rr,piv=Matrix([self.a[i]+B.a[i] for i in range(self.rows)]).rref()
        assert list(range(self.cols))==[c for c in piv if c<self.cols]
        assert all(c<self.cols for c in piv), 'inconsistent system'
        return Matrix([rr.a[i][self.cols:] for i in range(self.cols)]),None
    def inv(self):
        assert self.rows==self.cols
        I=Matrix([[int(i==j) for j in range(self.cols)] for i in range(self.rows)])
        return self.gauss_jordan_solve(I)[0]

def Rational(a,b=1):return F(a,b)
def zeros(r,c):return Matrix([[0]*c for _ in range(r)])
def ones(r,c):return Matrix([[1]*c for _ in range(r)])