如何在定义函数时将随机数数组传递给random.choice？

Question

我试图编写一个python代码来计算随机停止的和，但当随机生成的数字的大小超过5时会遇到问题

MemoryError: Unable to allocate 15.5 PiB for an array with shape (49, 53, 43, 46, 52, 53, 57, 52, 52) and data type float64

在以下代码中：

#Create a function tn.func to calculate randomly stopped sum
import numpy as np
#Define a function
def tn_fun(n):
    return sum(np.random.choice([50, 100, 200], n, replace=True, p=[0.3, 0.5, 0.2]))
N = np.random.poisson(50, 10)
# #Generate 10000 random values of N, using lambda = 50
TN = tn_fun(N)
print('Sample mean of the randomly stopped sum TN is',np.mean(TN))
print('Sample variance of the randomly stopped sum TN is', np.var(TN))

Answer 1

这就是你想要的

import numpy as np

# How many samples should there be?

# This is uniform between 5000 and 15000.
#N = np.random.randint(5000, 15000)

# This picks one number with Poisson distribution centered at 10000.
N = np.random.poisson(10000, 1)[0]

# Generate them.

TN = np.random.choice( [50,100,200], N, p=[0.3,0.5,0.2] )

print('Sample mean of the randomly stopped sum TN is',np.mean(TN))
print('Sample variance of the randomly stopped sum TN is', np.var(TN))