python中生成剔除部分数字的随机整数

要实现的功能

最开始是想简单地实现返回在区间[low, high)之间的随机整数，只是有一个值我不想他出现而已。如果这样的话，代码真的很简单。

def gen_randint(low, high, discard):
	result_list = list(range(low, high))
	result_list.remove(discard)
	np.random.shuffle(result_list)
	
	return result_list.pop()

但在做的过程中，想到可能会出现下面这几种情况。

1. 需要剔除的不是一个单值，而是一串数字
2. 可能需要返回一个随机数组
3. 可能只传了上界或者下界，那么这时候就要考虑到如何设置生成区间的问题

代码实现

# @Time    :2020/12/12 11:40
# @Author  :korolTung
# @FileName: user_function.py
import numpy as np

def randomint_plus(low, high=None, cutoff=None, size=None):
    """用来生成不包含一些中间数值的随机整数或序列

    Parameters
    ----------
    low : int
        生成随机整数的下界
    high: int
        生成随机整数的上界
    cutoff: int/list
        不需要包含的一个或多个值
    size: tuple
        需要产生的随机数规模
    Notes
    -----
    1. 在调用过程中，如果high, cutoff, size都缺省，就默认返回一个[0, low)的值
    2. 如果cutoff, size缺省，返回[low, high)的一个随机整数值
    3. 如果size缺省， 则返回一个[low, cutoff)U(cutoff, high)的随机整数值
    4. 返回一个给定size的矩阵，矩阵元素为[low, cutoff)U(cutoff, high)的随机整数值

    See Also
    --------
    np.random.randint()

    """
    if high is None:
        assert low is not None, "必须给定一个值作为上界"
        high = low
        low = 0
    number = 1  # 将size拉长成为一个向量
    if size is not None:
        for i in range(len(size)):
            number = number * size[i]

    if cutoff is None:  # 如果不需要剔除值，就通过numpy提供的函数生成随机整数
        random_result = np.random.randint(low, high, size=size)
    else:
        if number == 1:  # 返回一个随机整数
            random_result = randint_digit(low, high, cutoff)
        else:  # 返回一个形状为size的随机整数数组
            random_result = np.zeros(number, )
            for i in range(number):
                random_result[i] = randint_digit(low, high, cutoff)
            random_result = random_result.reshape(size)

    return random_result.astype(int)


def randint_digit(low, high, cutoff):
    """用来生成一个在区间[low, high)排除cutoff后的随机整数

    Parameters
    ----------
    low: int
        下限，能够取到
    high: int
        上限，不能够取到
    cutoff: int/list
        一个需要剔除的数或者数组，要求在(low, high)区间之间
    """
    digit_list = list(range(low, high))
    if type(cutoff) is int:  # 只需要剔除一个值
        if cutoff in digit_list:  # 如果需要剔除的值不存在，则不执行剔除操作
            digit_list.remove(cutoff)
    else:
        for i in cutoff:  # 需要剔除多个值的情况
            if i not in digit_list:  # 如果需要剔除的值不存在，则不执行剔除操作
                continue
            digit_list.remove(i)

    np.random.shuffle(digit_list)

    return digit_list.pop()  # 生成的序列打乱并且返回当前的随机值

if __name__ == '__main__':
    cutoff1 = 3
    cutoff2 = [3, 5, 10]
    low = 1
    high = 9
    result = randint_digit(low, high, cutoff2)

思路

思路很简单，但是要考虑到每种传入参数的可能，还是下了点功夫。总的来说就是我太菜了。好在有所收获。