时间戳

首页 > 解决方案 > 为 np.meshgrid 提供不同数量的输入参数

python - 为 np.meshgrid 提供不同数量的输入参数

问题描述

我有一个函数用于np.meshgrid获取所提供坐标的矩阵形式。我有一个参数dim来确定我正在使用的维度，并且需要返回一个dim维度沿轴 1 的数组。我在下面附加了一个 MWE。

import numpy as np

def func(dim=2):
    a = [1, 2]
    return np.array(np.meshgrid([a]*dim)).T.reshape(-1, dim)

func()
"""
returns
array([[1, 2],
       [1, 2]])
"""

但是我的预期输出是

array([[1, 1],
       [1, 2],
       [2, 1],
       [2, 2]])

，这是通过用 替换获得return np.array(np.meshgrid([a]*dim)).T.reshape(-1, dim)的return np.array(np.meshgrid(a, a)).T.reshape(-1, dim)。(a, a)请注意，我作为参数传入np.meshgrid, 因为dim=2. 如果dim=3，输入将是(a, a, a)等等。

我怎样才能做到这一点？如果任何其他功能可以做到这一点，我也对此持开放态度。谢谢。

标签： pythonarraysnumpy

解决方案

看看这些之间的区别：

def print_args(*args):
    print(args)

a = [1, 2]
print_args([a] * 2)
# ([[1, 2], [1, 2]],)

print_args(*[a] * 2)
# ([1, 2], [1, 2])

在第一个中，您正在传递一个列表列表。在第二个中，您正在解包（第一个）列表*，func因此每个子列表都是一个位置参数。

def func(dim=2):
    a = [1, 2]
    return np.array(np.meshgrid(*[a] * dim)).T.reshape(-1, dim)

func()
# array([[1, 1],
#        [1, 2],
#        [2, 1],
#        [2, 2]])

推荐阅读