Блог

Мне нужно создать функцию, которая может принимать неопределенное количество параметров (элементов для матрицы) и возвращать соответствующую квадратную матрицу. Я реализовал это, используя следующий подход.

 def square_matrix(size, *elements):

    numbers = list(elements)
    if size ** 2 != len(numbers):
        return "Number of elements does not match the size of the matrix"
    else:
        matrix = []
        factor = 0
        for i in range(0, size):
            row = []
            for j in range(factor * size, (factor   1) * size):
                row.append(numbers[j])
            factor  = 1
            matrix.append(row)
            i  = 1
        return matrix

print(square_matrix(3, 1, 2, 3, 4, 5, 6, 7, 8, 9))
# Output: [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
  

Хотя этот метод отлично работает для матриц меньшего размера, он кажется несколько неэффективным, поскольку использует вложенные циклы и кажется мне излишне длинным. Есть ли лучший / краткий способ реализовать то же самое?

можно ли просто использовать NumPy?

 import numpy as np

def square_matrix(size, *elements):
    return np.array(elements).reshape(size, size)
  

Для создания матрицы вы можете использовать понимание списка, которое более эффективно, чем для цикла.

Вы могли бы заменить циклы for в части else вашего оператора if приведенным ниже кодом.

 matrix = [[i for i in elements[j:j size]] for j in range(0,len(elements),size)]
  

Нет необходимости указывать размер в качестве аргумента. Если квадратный корень из длины аргументов является целым числом, то вы можете составить квадратную матрицу из элементов.

 import numpy as np
from math import sqrt

def square_matrix(*elements):
    size = sqrt(len(elements))
    if size.is_integer():
      return np.array(elements).reshape(int(size), int(size))
    else:
      raise RuntimeError("Number of elements is not sufficient to make a square matrix")


print(square_matrix(1, 2, 3, 4, 5, 6, 7, 8, 9))
# Output:
# array([[1, 2, 3],
#        [4, 5, 6],
#        [7, 8, 9]])