让函数支持整数与类数组输入的最Pythonic实现方案（Numpy）

问题背景 #

我正在开发一个用于DnD、Pathfinder等TRPG系统的整数分布计算工具，需要实现支持整数和类数组输入的PMF（概率质量函数）与CDF（累积分布函数）查询功能。

处理输入时，我用np.atleast_1d(np.asarray(N, dtype=int))确保输入至少为1维，满足内部逻辑要求，但这样会导致单个整数输入返回形状为(1,)的float数组，而非预期的0维float数组。目前我通过缓存输入形状、最后reshape的方式解决，但觉得这个方法太繁琐，想找Numpy内置的更高效实现方案。

我的最小可运行代码如下：

import numpy as np

class integerDistribution:

    def __init__(self, min_val, probability_distribution, rtol = 1e-10):
        """Initializes the instance from a minimum value and a finite list of probabilities.

        Args:
            min_val : int
                The minimum value of our distribution. As an example, rolling a d20 
                would have a minimum value of 1, while the sum of 2d6 has a minimum 
                value of 2.
            probability_distribution : array-like of floats
                Each entry at index 'idx' in the array corresponds to the probability 
                of obtaining the value 'idx + min_val' from our distribution.
            rtol: float
                When initializing, we perform a sanity check that the sum of 
                probabilities is 1 to within a tolerance of rtol. Defaults to 1e-10

        Raises:
            ValueError: If probabilities do not sum to 1 within the specified
                tolerance (rtol)      
        """
        if not np.isclose(np.sum(probability_distribution), 1.0, rtol):
            raise ValueError(f"Probabilities summed to {np.sum(probability_distribution)}")

        self.min_val = int(min_val)
        self.max_val = self.min_val + len(probability_distribution) - 1
        self.values = self.min_val + np.arange(len(probability_distribution), dtype = int)
        self.probability_distribution = np.array(probability_distribution)
        

    
    def pmf(self, N):
        """Evaluates the probability mass function (PMF) at given value(s)

        The probability mass function is defined as p(N) = P(X = N), 
        where X is a random variable sampled from 'self', and N is the value 
        for which the probability is being calculated. 

        Args:
            N : int or list of int
                The value(s) for which to compute the probability. 
                Can be a single integer or a list/array of integers.

        Returns:
            p(N) : float or list of float
                The probability of obtaining each value in N when sampled from 'self'. 
                The return type matches the input type: a single float if N is an 
                integer, or a list/array of floats if N is a list/array of integers.
        """
        N_shapecache = np.asarray(N, dtype = int)
        N = np.atleast_1d(N_shapecache)
        p_N = np.zeros(N.shape)
        valid_indices = np.where((self.min_val <= N)*(N <= self.max_val))
        p_N[valid_indices] = self.probability_distribution[N[valid_indices] - self.min_val]
        return p_N.reshape(N_shapecache.shape)

解决方案 #

方案1：标量与数组分支处理（推荐） #

直接判断输入是否为标量，分别处理两种情况，避免统一转成1维再reshape，逻辑清晰且性能最优：

def pmf(self, N):
    N_arr = np.asarray(N, dtype=int)
    # 处理标量输入
    if np.isscalar(N):
        if not (self.min_val <= N <= self.max_val):
            return 0.0
        return float(self.probability_distribution[N - self.min_val])
    # 处理数组输入
    p_N = np.zeros(N_arr.shape, dtype=float)
    valid_mask = (self.min_val <= N_arr) & (N_arr <= self.max_val)
    p_N[valid_mask] = self.probability_distribution[N_arr[valid_mask] - self.min_val]
    return p_N

方案2：简化形状还原逻辑 #

如果希望保持统一的数组处理逻辑，可以用ravel()替代np.atleast_1d，最后直接还原原始形状，写法更简洁：

def pmf(self, N):
    N_arr = np.asarray(N, dtype=int)
    N_1d = N_arr.ravel()  # 转成1维视图，比atleast_1d更高效
    p_N = np.zeros(N_1d.shape, dtype=float)
    valid_indices = np.where((self.min_val <= N_1d) & (N_1d <= self.max_val))
    p_N[valid_indices] = self.probability_distribution[N_1d[valid_indices] - self.min_val]
    return p_N.reshape(N_arr.shape)

方案3：用np.take简化赋值操作 #

利用np.take直接提取有效索引对应的概率值，结合zeros_like创建同形状数组，最后处理标量返回：

def pmf(self, N):
    N_arr = np.asarray(N, dtype=int)
    p_N = np.zeros_like(N_arr, dtype=float)
    valid_mask = (self.min_val <= N_arr) & (N_arr <= self.max_val)
    # 计算有效输入对应的分布索引
    indices = N_arr[valid_mask] - self.min_val
    p_N[valid_mask] = np.take(self.probability_distribution, indices)
    # 标量输入返回单个float，否则返回原形状数组
    return float(p_N) if N_arr.ndim == 0 else p_N

总结 #

• 方案1最直观，针对标量和数组做分支处理，减少了不必要的数组维度转换，性能最优；
• 方案2和3适合希望保持代码统一性的场景，写法更简洁，利用Numpy内置函数简化操作。

内容的提问来源于stack exchange，提问作者Minty