将一维 NumPy 数组重塑为接近正方形的二维数组

本文旨在解决将一维 NumPy 数组重塑为尽可能接近正方形的二维数组的问题。由于并非所有数组长度都能完美分解为两个相等的整数，因此我们需要找到两个因子，它们的乘积等于数组长度，并且这两个因子尽可能接近。本文将介绍两种实现该目标的 Python 函数，并提供详细的代码示例和解释。

寻找最佳的数组形状

当需要将一维 NumPy 数组重塑为二维数组时，我们通常希望得到的二维数组的形状尽可能接近正方形。这意味着我们需要找到两个整数 p 和 q，使得 p * q 等于数组的长度 n，并且 p 和 q 的值尽可能接近 sqrt(n)。

以下提供了两种方法来实现这个目标：

方法一：快速查找（适用于较小的 n）

这种方法通过迭代小于 sqrt(n) 的整数来寻找因子，并选择最接近 sqrt(n) 的因子。

import numpy as np
from math import isqrt
def np_squarishrt(n):
a = np.arange(1, isqrt(n) + 1, dtype=int)
b = n // a
i = np.where(a * b == n)[0][-1]
return a[i], b[i]

代码解释：

1. isqrt(n): 计算 n 的整数平方根。
2. np.arange(1, isqrt(n) + 1, dtype=int): 创建一个从 1 到 isqrt(n) 的整数数组。
3. b = n // a: 计算 n 除以 a 中每个元素的整数商。
4. np.where(a * b == n)[0][-1]: 找到 a * b == n 的索引，并选择最后一个索引。这是为了确保我们选择的是最接近 sqrt(n) 的因子。
5. return a[i], b[i]: 返回找到的两个因子。

使用示例：

a = np.arange(500)
b = a.reshape(np_squarishrt(len(a)))
print(b.shape)  # 输出: (20, 25)

方法二：通用方法（适用于较大的 n）

这种方法使用质因数分解和幂集来查找所有可能的因子组合，并选择最接近 sqrt(n) 的因子。

DeepSeek

幻方量化公司旗下的开源大模型平台

7087

查看详情

import numpy as np
from math import isqrt
from itertools import chain, combinations
def factors(n):
while n > 1:
for i in range(2, n + 1):
if n % i == 0:
n //= i
yield i
break
def uniq_powerset(iterable):
"""
Similar to powerset(it) but without repeats.
uniq_powerset([1,1,2]) --> (), (1,), (2,), (1, 1), (1, 2), (1, 1, 2)"""
s = list(iterable)
return chain.from_iterable(set(combinations(s, r)) for r in range(len(s)+1))
def squarishrt(n):
p = isqrt(n)
if p**2 == n:
return p, p
bestp = 1
f = list(factors(n))
for t in uniq_powerset(f):
if 2 * len(t) > len(f):
break
p = np.prod(t) if t else 1
q = n // p
if p > q:
p, q = q, p
if p > bestp:
bestp = p
return bestp, n // bestp

代码解释：

1. factors(n): 使用埃拉托斯特尼筛法找到 n 的所有质因数。
2. uniq_powerset(iterable): 生成一个幂集，但不包含重复的组合。
3. squarishrt(n):
   • 首先检查 n 是否是完全平方数。如果是，则返回 sqrt(n), sqrt(n)。
   • 否则，找到 n 的所有质因数，并生成所有可能的因子组合。
   • 对于每个因子组合，计算 p 和 q，并选择最接近 sqrt(n) 的因子。

使用示例：

a = np.arange(500)
b = a.reshape(squarishrt(len(a)))
print(b.shape)  # 输出: (20, 25)

注意事项：

• 方法二的计算复杂度较高，特别是对于较大的 n。
• 方法一对于较小的 n 来说通常足够快。
• 选择哪种方法取决于 n 的大小和性能要求。

总结：

本文介绍了两种将一维 NumPy 数组重塑为接近正方形的二维数组的方法。方法一适用于较小的 n，而方法二适用于较大的 n。选择哪种方法取决于 n 的大小和性能要求。通过使用这些方法，您可以轻松地将一维 NumPy 数组重塑为更易于处理的二维数组。