algorithm-in-python/math/numberTheory/modulo_equation.py

''' mbinary
#########################################################################
# File : modulo_equation.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.xyz
# Github: https://github.com/mbinary
# Created Time: 2018-3-4  21:14
# Description:
   `--` represents modulo symbol
#########################################################################
'''

import re

from gcd import xgcd
from euler import phi
from isPrime import isPrime
from factor import factor


def ind(m, g):
    ''' mod m ,primary root g  ->  {n:indg n}'''
    return {j: i for i in range(m-1)
            for j in range(m) if (g**i-j) % m == 0}


def gs(m, num=100):
    '''return list of  m's  primary roots below num'''
    p = phi(m)
    mp = factor(p)
    checkLst = [p//i for i in mp]
    return [i for i in range(2, num) if all((i**n-1) % m != 0 for n in checkLst)]


def minG(m):
    p = phi(m)
    mp = factor(p)
    checkLst = [p//i for i in mp]
    i = 2
    while 1:
        if all((i**n-1) % m != 0 for n in checkLst):
            return i
        i += 1


class solve:
    def __init__(self, equ=None):
        self.linearPat = re.compile(r'\s*(\d+)\s*--\s*(\d+)[\s\(]*mod\s*(\d+)')
        self.sol = []
        #self.m = m
        #self.ind_mp = ind(m,minG(m))

    def noSol(self):
        print('equation {equ} has no solution'.format(equ=self.equ))

    def error(self):
        print("Error! The divisor m must be postive integer")

    def solveLinear(self, a, b, m):
        '''ax--b(mod m): solve linear equation with one unknown
            return  ([x1,x2,...],m)
        '''
        a, b, m = self.check(a, b, m)
        g, x, y = xgcd(a, m)
        if a*b % g != 0:
            self.noSol()
            return None
        sol = x*b//g
        m0 = m//g
        sols = [(sol+i*m0) % m for i in range(g)]
        print('{}x--{}(mod {}), solution: {} mod {}'.format(a, b, m, sols, m))
        return (sols, m)

    def check(self, a, b, m):
        if m <= 0:
            self.error()
            return None
        if a < 0:
            a, b = -a, -b  # important
        if b < 0:
            b += -b//m * m
        return a, b, m

    def solveHigh(self, a, n, b, m):
        ''' ax^n -- b (mod m)  ind_mp is a dict of  m's {n: indg n}'''
        ind_mp = ind(m, minG(m))
        tmp = ind_mp[b] - ind_mp[a]
        if tmp < 0:
            tmp += m
        sol = self.solveLinear(n, tmp, phi(m))
        re_mp = {j: i for i, j in ind_mp.items()}
        sols = [re_mp[i] for i in sol[0]]
        print('{}x^{}--{}(mod {}),  solution: {} mod {}'.format(a, n, b, m, sols, m))
        return sols, m

    def solveGroup(self, tups):
        '''tups is a list of tongyu equation groups, like
            [(a1,b1,m1),(a2,b2,m2)...]
            and, m1,m2... are all primes
        '''
        mp = {}
        print('solving group of equations: ')
        for a, b, m in tups:
            print('{}x--{}(mod {})'.format(a, b, m))
            if m in mp:
                if mp[m][0]*b != mp[m][1]*a:
                    self.noSol()
                    return
            else:
                mp[m] = (a, b)
        product = 1
        for i in mp.keys():
            product *= i
        sol = [0]
        for i in mp:
            xs, m = self.solveLinear(product//i*mp[i][0], 1, i)
            new = []
            for x in xs:
                cur = x*product//i*mp[i][1]
                for old in sol:
                    new.append(old+cur)
            sol = new
        sol = [i % product for i in sol]
        print('final solution: {} mod {}'.format(sol, product))
        return sol, product

    def __call__(self):
        s = input('输入同余方程，用--代表同于号，形如3--5(mod 7)代表3x模7同余于5')
        li = self.linearPat.findall(s)
        li = [(int(a), int(b), int(m)) for a, b, m in li]
        print(self.solveLinear(li[0]))


if __name__ == '__main__':
    solver = solve()
    res = solver.solveLinear(3, 6, 9)
    print()
    res = solver.solveHigh(1, 8, 3, 11)
    print()
    res = solver.solveGroup([(5, 11, 2), (3, 8, 5), (4, 1, 7)])