ranea / qipy Goto Github PK

w = Symbol('w')
W = (-1 + sqrt(3)*I)/2 # or if you prefer exp(2*pi*I/3)

def rewrite_in_W(x):
    x = x.subs(w, W)
    a, b3 = re(x), im(x)/sqrt(3)
    return a+b3, 2*b3

def divmod(x, y, *, round=floor):
    qa, qb = rewrite_in_W(x/y)
    qa, qb = round(qa), round(qb)
    q = qa + qb*w
    r = x - q*y
    ra, rb = rewrite_in_W(x - q*y)
    assert (ra, rb) == (round(ra), round(rb))
    r = ra + rb*w
    return q, r

assert divmod(-7+7*w, 1+6*w) == (2+w, -3)
assert divmod(-7+7*w, 1+6*w, round=ceil) == (3+2*w, 2-w)

Ei = QuadraticIntegerRing(-3)
w = (-1 + sqrt(3)*I)/2

def rewrite_in_w(x):
    a, b3 = re(x), im(x)/sqrt(3)
    return a+b3, 2*b3

def divmod(x, y, *, round=floor):
    qa, qb = rewrite_in_w(x.value/y.value)
    qa, qb = round(qa), round(qb)
    q = qa + qb*w
    r = x - q*y
    return type(x)(q), type(x)(r)

ei77 = Ei(-7+7*w)
ei16 = Ei(1+6*w)

assert divmod(ei77, ei16) == (Ei(2+w), Ei(-3))
assert divmod(ei77, ei16, round=ceil) == (Ei(3+2*w), Ei(2-w))