In addition to creating a helper method that receives the operator as a parameter, as shown in the other responses, you can use other magic methods to simplify the logic of your class.
Consider the start class:
class Matrix:
def __init__(self, matrix):
self.matrix = matrix
You can make your object iterable by setting the __iter__
method:
class Matrix:
def __init__(self, matrix):
self.matrix = matrix
def __iter__(self):
yield from iter(self.matrix)
In this way, you can iterate over the rows of the array with for
:
>>> m = Matrix([[1, 2, 3], [4, 5, 6]])
>>> for linha in m:
... print(linha)
[1, 2, 3]
[4, 5, 6]
You can use the __len__
method to return the array dimension:
class Matrix:
def __init__(self, matrix):
self.matrix = matrix
def __iter__(self):
yield from iter(self.matrix)
def __len__(self):
return len(self.matrix)
And so, do:
>>> print(len(m))
2
You can implement the __getitem__
method to facilitate access to array rows:
class Matrix:
def __init__(self, matrix):
self.matrix = matrix
def __iter__(self):
yield from iter(self.matrix)
def __len__(self):
return len(self.matrix)
def __getitem__(self, key):
return self.matrix[key]
Can do:
>>> print(m[0])
[1, 2, 3]
Finally, implementing the addition and subtraction methods, your class could look like:
class Matrix:
def __init__(self, matrix):
self.matrix = matrix
@property
def size(self):
return (len(self), len(self[0]))
def __iter__(self):
yield from iter(self.matrix)
def __len__(self):
return len(self.matrix)
def __getitem__(self, key):
return self.matrix[key]
def __add__(self, other):
if self.size != other.size:
raise Exception('Matrizes de tamanhos diferentes')
return Matrix(
[
[self[i][j] + other[i][j] for j in range(len(self[i]))]
for i in range(len(self))
]
)
def __sub__(self, other):
if self.size != other.size:
raise Exception('Matrizes de tamanhos diferentes')
return Matrix(
[
[self[i][j] - other[i][j] for j in range(len(self[i]))]
for i in range(len(self))
]
)