1- from itertools import chain
2- from typing import Dict , Iterable , List , Optional , Tuple
3-
4-
51def is_prime (n : int ) -> bool :
62 """Primality test using 6k+-1 optimization."""
73 if n <= 3 :
@@ -14,147 +10,3 @@ def is_prime(n: int) -> bool:
1410 return False
1511 i += 6
1612 return True
17-
18-
19- class Matrix :
20- """
21- Creates a square matrix of the dimensions (size x size).
22- Arguments:
23- `size:` The size of the `n x n` matrix.
24- `binary:` If set to True then all calculations are done over GF(2).
25- `elements:` The function from pos to value at pos. `f(pos) = value`.
26- """
27-
28- def __init__ (
29- self ,
30- size : int ,
31- elements : Optional [Dict [Tuple [int , int ], int ]] = None ,
32- binary : bool = True ,
33- ) -> None :
34- if elements :
35- self .elements : Dict [Tuple [int , int ], int ] = {
36- key : (1 if binary and val else val ) for key , val in elements .items ()
37- }
38- else :
39- self .elements = {}
40- self .binary = binary
41- self ._size = size
42-
43- def __getitem__ (self , pos ) -> int :
44- """Get value in row, col specified"""
45- return self .elements [pos ] if pos in self .elements else 0
46-
47- def __setitem__ (self , pos , value ) -> None :
48- """Sets position in the matrix to val"""
49- assert 0 <= pos [0 ] < len (self )
50- assert 0 <= pos [1 ] < len (self )
51- self .elements [pos ] = value
52-
53- def __add__ (self , other : "Matrix" ) -> "Matrix" :
54- """Adds two matrices together. Both classes have to be binary or not binary"""
55- if not isinstance (other , Matrix ):
56- return NotImplemented
57- if self .binary != other .binary :
58- return NotImplemented
59- assert len (self ) == len (other )
60- new_mat = Matrix (size = len (self ))
61- funcs = [self .elements .items (), other .elements .items ()]
62- if self .binary :
63- for idx , val in chain (* funcs ):
64- new_mat [idx ] ^= val
65- else :
66- for idx , val in chain (* funcs ):
67- new_mat [idx ] += val
68- return new_mat
69-
70- def __mul__ (self , other : "Matrix" ) -> "Matrix" :
71- if not isinstance (other , Matrix ):
72- return NotImplemented
73- if self .binary != other .binary :
74- return NotImplemented
75- assert len (self ) == len (other )
76- mat = Matrix (size = len (self ), elements = {})
77- for row in range (len (self )):
78- for col in range (len (self )):
79- for k in range (len (self )):
80- mat [(row , col )] += self [(row , k )] * other [(k , col )]
81- return mat
82-
83- def __len__ (self ) -> int :
84- return self ._size
85-
86- def __str__ (self ) -> str :
87- string = ""
88- for row in range (len (self )):
89- for col in range (len (self )):
90- cell = self [row , col ]
91- string += f"|{ cell }
92- string += "|\n "
93- return string
94-
95- def __repr__ (self ) -> str :
96- return f"Matrix({ len (self )} { self .elements }
97-
98- def __eq__ (self , other : object ) -> bool :
99- """Returns True if both matrices are equal."""
100- if not isinstance (other , Matrix ):
101- return NotImplemented
102- if self .binary != other .binary :
103- return NotImplemented
104- if len (self ) != len (other ):
105- return False
106- for row in range (len (self )):
107- for col in range (len (self )):
108- if self [row , col ] != other [row , col ]:
109- return False
110- return True
111-
112- def test_sum (self ) -> int :
113- """Sums up all the values in the matrix."""
114- return sum (self .elements .values ())
115-
116- def hamming_sums (self , row : bool = True ) -> List [int ]:
117- """Gives the hamming sums for either the rows or columns"""
118- key_idx = 0 if row else 1
119- out = []
120- for idx in range (len (self )):
121- out .append (0 )
122- for key , val in self .elements .items ():
123- if key [key_idx ] == idx :
124- out [idx ] += val
125- return out
126-
127- def row_labels (self ) -> Tuple [int , ...]:
128- """
129- Looks at each row as a binary number and returns the binary value for each.
130- """
131- out_list : List [int ] = []
132- for row_idx in range (len (self )):
133- num = 0
134- for key , val in self .elements .items ():
135- if key [0 ] == row_idx and val :
136- num += 2 ** (key [1 ])
137- out_list .append (num )
138- return tuple (sorted (out_list ))
139-
140- def apply_perm (self , perm : Iterable [int ], row : bool = True ) -> "Matrix" :
141- """Returns a new matrix which has the given permutation
142- applied on either the rows or columns of the matrix"""
143- new_func : Dict [Tuple [int , int ], int ] = {}
144- for n_idx , p_idx in enumerate (perm ):
145- for idx in range (len (self )):
146- if row :
147- pos = (p_idx , idx )
148- else :
149- pos = (idx , p_idx )
150-
151- if pos in self .elements :
152- if row :
153- new_func [n_idx , idx ] = self [pos ]
154- else :
155- new_func [idx , n_idx ] = self [pos ]
156-
157- return Matrix (len (self ), elements = new_func )
158-
159- def __hash__ (self ) -> int :
160- return hash (str (self ))
0 commit comments