cp-library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub kobejean/cp-library

:warning: cp_library/math/linalg/mat/mod/mat_cls.py

Depends on

Code

import cp_library.__header__
from typing import Container, Sequence
from numbers import Number
from cp_library.io.parser_cls import Parsable, Parser, TokenStream
import cp_library.math.__header__
import cp_library.math.linalg.__header__
from cp_library.math.linalg.elm_wise_in_place_mixin import ElmWiseInPlaceMixin
import cp_library.math.linalg.mat.__header__

class Mat(Parsable, Container, ElmWiseInPlaceMixin):

    def __init__(self, flat, N, M):
        self.data, self.N, self.M = flat, N, M

    def elm_wise(self, other, op):
        cls = type(self)
        if isinstance(other, Number):
            return cls([op(elm, other) for elm in self.data])
        if isinstance(other, Sequence):
            return cls([op(self.data[i], elm) for i, elm in enumerate(other)])
        raise ValueError("Operand must be a number or a tuple of the same length")
    
    def ielm_wise(self, other, op):
        data = self.data
        if isinstance(other, Number):
            for i in range(len(data)):
                data[i] = op(data[i], other)
        elif isinstance(other, Sequence) and len(data) == len(other):
            for i, elm in enumerate(other):
                data[i] = op(data[i], elm)
        else:
            raise ValueError("Operand must be a number or a list of the same length")
        return self

    def __len__(self):
        return self.N

    def __getitem__(self, ij: tuple[int, int]):
        i, j = ij
        return self.data[i*self.M+j]
    
    def __setitem__(self, ij: tuple[int, int], val: int):
        i, j = ij
        self.data[i*self.M+j] = val
    
    def __contains__(self, x: int) -> bool:
        return x in self.data

    def __matmul__(A,B):
        assert A.M == B.N, f"Dimension mismatch {A.M = } {B.N = }"
        N,M = A.N, B.M
        cls = type(A)
        R = cls([0]*(M*N))
        for irow in range(0,N*M,M):
            for k in range(A.M):
                krow, a = k*M, A.data[irow+k]
                for j in range(M):
                    R.data[irow+j] = B.data[krow+j]*a + R.data[irow+j]
        return R
    
    def __pow__(A,K):
        R = A.copy() if K & 1 else type(A).identity(A.N)
        for i in range(1,K.bit_length()):
            A = A @ A
            if K >> i & 1:
                R = R @ A
        return R 

    @classmethod
    def identity(cls, N):
        data = [0]*(N*N)
        for i in range(0,N*N,N+1): data[i] = 1
        return cls(data)
    
    def copy(self):
        cls = type(self)
        obj = cls.__new__(cls)
        obj.N, obj.M, obj.data = self.N, self.M, self.data
        return obj
    
    @classmethod
    def compile(cls, N: int, M: int, T: type = int):
        elm, size = Parser.compile(T), N*M
        def parse(ts: TokenStream):
            return cls([elm(ts) for _ in range(size)])
        return parse
    
    def __repr__(self) -> str:
        return '\n'.join(' '.join(str(elm) for elm in row) for row in self)

from cp_library.math.mod.mint_cls import mint
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Container, Sequence
from numbers import Number

import typing
from collections import deque
from types import GenericAlias 
from typing import Callable, Collection, Iterator, Union
import os
import sys
from io import BytesIO, IOBase


class FastIO(IOBase):
    BUFSIZE = 8192
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = "x" in file.mode or "r" not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        BUFSIZE = self.BUFSIZE
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        BUFSIZE = self.BUFSIZE
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b"\n") + (not b)
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            self.buffer.truncate(0), self.buffer.seek(0)


class IOWrapper(IOBase):
    stdin: 'IOWrapper' = None
    stdout: 'IOWrapper' = None
    
    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable

    def write(self, s):
        return self.buffer.write(s.encode("ascii"))
    
    def read(self):
        return self.buffer.read().decode("ascii")
    
    def readline(self):
        return self.buffer.readline().decode("ascii")

sys.stdin = IOWrapper.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper.stdout = IOWrapper(sys.stdout)
from typing import TypeVar
_T = TypeVar('T')

class TokenStream(Iterator):
    stream = IOWrapper.stdin

    def __init__(self):
        self.queue = deque()

    def __next__(self):
        if not self.queue: self.queue.extend(self._line())
        return self.queue.popleft()
    
    def wait(self):
        if not self.queue: self.queue.extend(self._line())
        while self.queue: yield
 
    def _line(self):
        return TokenStream.stream.readline().split()

    def line(self):
        if self.queue:
            A = list(self.queue)
            self.queue.clear()
            return A
        return self._line()
TokenStream.default = TokenStream()

class CharStream(TokenStream):
    def _line(self):
        return TokenStream.stream.readline().rstrip()
CharStream.default = CharStream()


ParseFn = Callable[[TokenStream],_T]
class Parser:
    def __init__(self, spec: Union[type[_T],_T]):
        self.parse = Parser.compile(spec)

    def __call__(self, ts: TokenStream) -> _T:
        return self.parse(ts)
    
    @staticmethod
    def compile_type(cls: type[_T], args = ()) -> _T:
        if issubclass(cls, Parsable):
            return cls.compile(*args)
        elif issubclass(cls, (Number, str)):
            def parse(ts: TokenStream): return cls(next(ts))              
            return parse
        elif issubclass(cls, tuple):
            return Parser.compile_tuple(cls, args)
        elif issubclass(cls, Collection):
            return Parser.compile_collection(cls, args)
        elif callable(cls):
            def parse(ts: TokenStream):
                return cls(next(ts))              
            return parse
        else:
            raise NotImplementedError()
    
    @staticmethod
    def compile(spec: Union[type[_T],_T]=int) -> ParseFn[_T]:
        if isinstance(spec, (type, GenericAlias)):
            cls = typing.get_origin(spec) or spec
            args = typing.get_args(spec) or tuple()
            return Parser.compile_type(cls, args)
        elif isinstance(offset := spec, Number): 
            cls = type(spec)  
            def parse(ts: TokenStream): return cls(next(ts)) + offset
            return parse
        elif isinstance(args := spec, tuple):      
            return Parser.compile_tuple(type(spec), args)
        elif isinstance(args := spec, Collection):  
            return Parser.compile_collection(type(spec), args)
        elif isinstance(fn := spec, Callable): 
            def parse(ts: TokenStream): return fn(next(ts))
            return parse
        else:
            raise NotImplementedError()

    @staticmethod
    def compile_line(cls: _T, spec=int) -> ParseFn[_T]:
        if spec is int:
            fn = Parser.compile(spec)
            def parse(ts: TokenStream): return cls([int(token) for token in ts.line()])
            return parse
        else:
            fn = Parser.compile(spec)
            def parse(ts: TokenStream): return cls([fn(ts) for _ in ts.wait()])
            return parse

    @staticmethod
    def compile_repeat(cls: _T, spec, N) -> ParseFn[_T]:
        fn = Parser.compile(spec)
        def parse(ts: TokenStream): return cls([fn(ts) for _ in range(N)])
        return parse

    @staticmethod
    def compile_children(cls: _T, specs) -> ParseFn[_T]:
        fns = tuple((Parser.compile(spec) for spec in specs))
        def parse(ts: TokenStream): return cls([fn(ts) for fn in fns])  
        return parse
            
    @staticmethod
    def compile_tuple(cls: type[_T], specs) -> ParseFn[_T]:
        if isinstance(specs, (tuple,list)) and len(specs) == 2 and specs[1] is ...:
            return Parser.compile_line(cls, specs[0])
        else:
            return Parser.compile_children(cls, specs)

    @staticmethod
    def compile_collection(cls, specs):
        if not specs or len(specs) == 1 or isinstance(specs, set):
            return Parser.compile_line(cls, *specs)
        elif (isinstance(specs, (tuple,list)) and len(specs) == 2 and isinstance(specs[1], int)):
            return Parser.compile_repeat(cls, specs[0], specs[1])
        else:
            raise NotImplementedError()

class Parsable:
    @classmethod
    def compile(cls):
        def parser(ts: TokenStream): return cls(next(ts))
        return parser



import operator
from math import hypot


class ElmWiseMixin:
    def elm_wise(self, other, op):
        if isinstance(other, Number):
            return type(self)(op(x, other) for x in self)
        if isinstance(other, Sequence):
            return type(self)(op(x, y) for x, y in zip(self, other))
        raise ValueError("Operand must be a number or a tuple of the same length")

    def __add__(self, other): return self.elm_wise(other, operator.add)
    def __radd__(self, other): return self.elm_wise(other, operator.add)
    def __sub__(self, other): return self.elm_wise(other, operator.sub)
    def __rsub__(self, other): return self.elm_wise(other, lambda x,y: operator.sub(y,x))
    def __mul__(self, other): return self.elm_wise(other, operator.mul)
    def __rmul__(self, other): return self.elm_wise(other, operator.mul)
    def __truediv__(self, other): return self.elm_wise(other, operator.truediv)
    def __rtruediv__(self, other): return self.elm_wise(other, lambda x,y: operator.truediv(y,x))
    def __floordiv__(self, other): return self.elm_wise(other, operator.floordiv)
    def __rfloordiv__(self, other): return self.elm_wise(other, lambda x,y: operator.floordiv(y,x))
    def __mod__(self, other): return self.elm_wise(other, operator.mod)

    def distance(self: 'ElmWiseMixin', other: 'ElmWiseMixin'):
        diff = other-self
        return hypot(*diff)
    
    def magnitude(vec: 'ElmWiseMixin'):
        return hypot(*vec)
    
    def norm(vec: 'ElmWiseMixin'):
        return vec / vec.magnitude()

class ElmWiseInPlaceMixin(ElmWiseMixin):
    def ielm_wise(self, other, op):
        if isinstance(other, Number):
            for i in range(len(self)):
                self[i] = op(self[i], other)
        elif isinstance(other, Sequence) and len(self) == len(other):
            for i in range(len(self)):
                self[i] = op(self[i], other[i])
        else:
            raise ValueError("Operand must be a number or a list of the same length")
        return self
    
    def __iadd__(self, other): return self.ielm_wise(other, operator.add)
    def __isub__(self, other): return self.ielm_wise(other, operator.sub)
    def __imul__(self, other): return self.ielm_wise(other, operator.mul)
    def __itruediv__(self, other): return self.ielm_wise(other, operator.truediv)
    def __ifloordiv__(self, other): return self.ielm_wise(other, operator.floordiv)
    def __imod__(self, other): return self.ielm_wise(other, operator.mod)


class Mat(Parsable, Container, ElmWiseInPlaceMixin):

    def __init__(self, flat, N, M):
        self.data, self.N, self.M = flat, N, M

    def elm_wise(self, other, op):
        cls = type(self)
        if isinstance(other, Number):
            return cls([op(elm, other) for elm in self.data])
        if isinstance(other, Sequence):
            return cls([op(self.data[i], elm) for i, elm in enumerate(other)])
        raise ValueError("Operand must be a number or a tuple of the same length")
    
    def ielm_wise(self, other, op):
        data = self.data
        if isinstance(other, Number):
            for i in range(len(data)):
                data[i] = op(data[i], other)
        elif isinstance(other, Sequence) and len(data) == len(other):
            for i, elm in enumerate(other):
                data[i] = op(data[i], elm)
        else:
            raise ValueError("Operand must be a number or a list of the same length")
        return self

    def __len__(self):
        return self.N

    def __getitem__(self, ij: tuple[int, int]):
        i, j = ij
        return self.data[i*self.M+j]
    
    def __setitem__(self, ij: tuple[int, int], val: int):
        i, j = ij
        self.data[i*self.M+j] = val
    
    def __contains__(self, x: int) -> bool:
        return x in self.data

    def __matmul__(A,B):
        assert A.M == B.N, f"Dimension mismatch {A.M = } {B.N = }"
        N,M = A.N, B.M
        cls = type(A)
        R = cls([0]*(M*N))
        for irow in range(0,N*M,M):
            for k in range(A.M):
                krow, a = k*M, A.data[irow+k]
                for j in range(M):
                    R.data[irow+j] = B.data[krow+j]*a + R.data[irow+j]
        return R
    
    def __pow__(A,K):
        R = A.copy() if K & 1 else type(A).identity(A.N)
        for i in range(1,K.bit_length()):
            A = A @ A
            if K >> i & 1:
                R = R @ A
        return R 

    @classmethod
    def identity(cls, N):
        data = [0]*(N*N)
        for i in range(0,N*N,N+1): data[i] = 1
        return cls(data)
    
    def copy(self):
        cls = type(self)
        obj = cls.__new__(cls)
        obj.N, obj.M, obj.data = self.N, self.M, self.data
        return obj
    
    @classmethod
    def compile(cls, N: int, M: int, T: type = int):
        elm, size = Parser.compile(T), N*M
        def parse(ts: TokenStream):
            return cls([elm(ts) for _ in range(size)])
        return parse
    
    def __repr__(self) -> str:
        return '\n'.join(' '.join(str(elm) for elm in row) for row in self)


    
class mint(int):
    mod: int
    zero: 'mint'
    one: 'mint'
    two: 'mint'
    cache: list['mint']

    def __new__(cls, *args, **kwargs):
        if 0<= (x := int(*args, **kwargs)) <= 2:
            return cls.cache[x]
        else:
            return cls.fix(x)

    @classmethod
    def set_mod(cls, mod: int):
        mint.mod = cls.mod = mod
        mint.zero = cls.zero = cls.cast(0)
        mint.one = cls.one = cls.fix(1)
        mint.two = cls.two = cls.fix(2)
        mint.cache = cls.cache = [cls.zero, cls.one, cls.two]

    @classmethod
    def fix(cls, x): return cls.cast(x%cls.mod)

    @classmethod
    def cast(cls, x): return super().__new__(cls,x)

    @classmethod
    def mod_inv(cls, x):
        a,b,s,t = int(x), cls.mod, 1, 0
        while b: a,b,s,t = b,a%b,t,s-a//b*t
        if a == 1: return cls.fix(s)
        raise ValueError(f"{x} is not invertible in mod {cls.mod}")
    
    @property
    def inv(self): return mint.mod_inv(self)

    def __add__(self, x): return mint.fix(super().__add__(x))
    def __radd__(self, x): return mint.fix(super().__radd__(x))
    def __sub__(self, x): return mint.fix(super().__sub__(x))
    def __rsub__(self, x): return mint.fix(super().__rsub__(x))
    def __mul__(self, x): return mint.fix(super().__mul__(x))
    def __rmul__(self, x): return mint.fix(super().__rmul__(x))
    def __floordiv__(self, x): return self * mint.mod_inv(x)
    def __rfloordiv__(self, x): return self.inv * x
    def __truediv__(self, x): return self * mint.mod_inv(x)
    def __rtruediv__(self, x): return self.inv * x
    def __pow__(self, x): 
        return self.cast(super().__pow__(x, self.mod))
    def __neg__(self): return mint.mod-self
    def __pos__(self): return self
    def __abs__(self): return self
Back to top page