Skip to content
16 changes: 8 additions & 8 deletions maths/matrix_exponentiation.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,15 @@


class Matrix:
def __init__(self, arg):
def __init__(self, arg: list[list] | int) -> None:
if isinstance(arg, list): # Initializes a matrix identical to the one provided.
self.t = arg
self.n = len(arg)
else: # Initializes a square matrix of the given size and set values to zero.
self.n = arg
self.t = [[0 for _ in range(self.n)] for _ in range(self.n)]

def __mul__(self, b):
def __mul__(self, b: Matrix) -> Matrix:
matrix = Matrix(self.n)
for i in range(self.n):
for j in range(self.n):
Expand All @@ -28,7 +28,7 @@ def __mul__(self, b):
return matrix


def modular_exponentiation(a, b):
def modular_exponentiation(a: Matrix, b: int) -> Matrix:
matrix = Matrix([[1, 0], [0, 1]])
while b > 0:
if b & 1:
Expand All @@ -38,7 +38,7 @@ def modular_exponentiation(a, b):
return matrix


def fibonacci_with_matrix_exponentiation(n, f1, f2):
def fibonacci_with_matrix_exponentiation(n: int, f1: int, f2: int) -> int:
"""
Returns the nth number of the Fibonacci sequence that
starts with f1 and f2
Expand All @@ -64,7 +64,7 @@ def fibonacci_with_matrix_exponentiation(n, f1, f2):
return f2 * matrix.t[0][0] + f1 * matrix.t[0][1]


def simple_fibonacci(n, f1, f2):
def simple_fibonacci(n: int, f1: int, f2: int) -> int:
"""
Returns the nth number of the Fibonacci sequence that
starts with f1 and f2
Expand Down Expand Up @@ -95,7 +95,7 @@ def simple_fibonacci(n, f1, f2):
return f2


def matrix_exponentiation_time():
def matrix_exponentiation_time() -> float:
setup = """
from random import randint
from __main__ import fibonacci_with_matrix_exponentiation
Expand All @@ -106,7 +106,7 @@ def matrix_exponentiation_time():
return exec_time


def simple_fibonacci_time():
def simple_fibonacci_time() -> float:
setup = """
from random import randint
from __main__ import simple_fibonacci
Expand All @@ -119,7 +119,7 @@ def simple_fibonacci_time():
return exec_time


def main():
def main() -> None:
matrix_exponentiation_time()
simple_fibonacci_time()

Expand Down