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

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]

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

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

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