python - NumPy 相当于 Matlab 的 magic()

Question

在 Ocatave/Matlab 中，我可以使用 magic()得到一个幻方，例如，

magic(4)

  16    2    3   13
   5   11   10    8
   9    7    6   12
   4   14   15    1

定义:幻方是一个 N×N 的数字网格，其中每一行、每一列和主对角线的元素总和为相同的数字(等于 N(N^2+1)/2)。

如何使用 NumPy 生成相同的内容？

Answer 1

这个实现遵循 Matlab 的，应该给出完全相同的结果，但有以下异常(exception):如果 n < 3，它会抛出一个错误，而不是返回一个非幻方 [[1, 3], [4, 2] ] 当 n=2 时，就像 Matlab 那样。

照例分为三种情况:奇数、能被4整除、偶数不能被4整除，最后一种最复杂。

def magic(n):
  n = int(n)
  if n < 3:
    raise ValueError("Size must be at least 3")
  if n % 2 == 1:
    p = np.arange(1, n+1)
    return n*np.mod(p[:, None] + p - (n+3)//2, n) + np.mod(p[:, None] + 2*p-2, n) + 1
  elif n % 4 == 0:
    J = np.mod(np.arange(1, n+1), 4) // 2
    K = J[:, None] == J
    M = np.arange(1, n*n+1, n)[:, None] + np.arange(n)
    M[K] = n*n + 1 - M[K]
  else:
    p = n//2
    M = magic(p)
    M = np.block([[M, M+2*p*p], [M+3*p*p, M+p*p]])
    i = np.arange(p)
    k = (n-2)//4
    j = np.concatenate((np.arange(k), np.arange(n-k+1, n)))
    M[np.ix_(np.concatenate((i, i+p)), j)] = M[np.ix_(np.concatenate((i+p, i)), j)]
    M[np.ix_([k, k+p], [0, k])] = M[np.ix_([k+p, k], [0, k])]
  return M 

我还写了一个函数来测试这个:

def test_magic(ms):
  n = ms.shape[0]
  s = n*(n**2+1)//2 
  columns = np.all(ms.sum(axis=0) == s)
  rows = np.all(ms.sum(axis=1) == s)
  diag1 = np.diag(ms).sum() == s 
  diag2 = np.diag(ms[::-1, :]).sum() == s
  return columns and rows and diag1 and diag2 

尝试 [test_magic(magic(n)) for n in range(3, 20)] 检查正确性。