Coverage for src/pqlattice/integer/_integer.py: 96%

1 from typing import cast, overload

2

3 from .._utils import as_integer

4 from ..typing import Array

5

6

7@overload

8 def eea(a: int, b: int) -> tuple[int, int, int]: ...

9

10

11@overload

12 def eea(a: Array, b: int) -> tuple[Array, Array, Array]: ...

13

14

15 def eea(a: int | Array, b: int) -> tuple[int, int, int] | tuple[Array, Array, Array]:

16 """_summary_

17

18 Parameters

19 ----------

20 a : int | Array

21 _description_

22 b : int

23 _description_

24

25 Returns

26 -------

27 tuple[int, int, int] | tuple[Array, Array, Array]

28 _description_

29 """

30 if isinstance(a, int):

31 return _eea(a, b)

32 else:

33 return tuple(as_integer([_eea(cast(int, el), b) for el in a]).T)

34

35

36 def _eea(a: int, b: int) -> tuple[int, int, int]:

37 if a == 0 and b == 0:

38 raise ValueError("a and b can't be both zero")

39

40 old_s, s = 1, 0

41 old_r, r = a, b

42 while r != 0:

43 q = old_r // r

44 old_r, r = r, old_r - q * r

45 old_s, s = s, old_s - q * s

46

47 t = 0 if b == 0 else (old_r - old_s * a) // b

48 s = old_s

49 gcd = old_r

50 if gcd < 0:

51 gcd = -gcd

52 s = -s

53 t = -t

54

55 return gcd, s, t