Python 要素数の和が一定値となる制約の元に、整数の乱数を作成する「numpy」

'22/03/12更新：自然な分布になる雛形コードも追記しました。
　本記事では、例えば複数のパラメータA, B, Cの3つがあって、それらの和が一定値Sの制約条件の元に、整数の乱数を作成する雛形コードを載せました。
下図は、S＝150と指定して2000水準作成した例です。生成したデータはcsvファイルに出力する仕様です。

▼1. 各要素の生成する乱数の範囲を確率で決定する場合
下図のように、A, B, Cの分布はおよそ同じ分布で生成されます。和が一定の場合はこれが自然な分布と想像できます。

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import random
import math
import pandas as pd

# パラメータ
parameter_name_list = ['A', 'B', 'C']
S = 150 # 制約条件 S＝A＋B＋C
N = 2000 # 水準数

data_list = [] # データを格納するリスト

i = 0
while i < N:
    # Aの値を確率によって決める
    flag_A = np.random.choice(['low', 'center', 'high'], p=[1/3, 1/3, 1/3])
    if flag_A == 'low': # この例では、Aは0~50の範囲で乱数により決定
        A = np.random.randint(0, S / 3)
    elif flag_A == 'center': # この例では、Aは51~100の範囲で乱数により決定
        A = random.randint(S / 3, S / 3 * 2)
    else: # この例では、Aは101~150の範囲で乱数により決定
        A = random.randint(S / 3 * 2, S)
    #pint(flag_A, A)
    
    if S == A or A == 0: # 0割は飛ばして、次のループへ
        continue
    else:
        # Bの値を確率によって決める
        flag_B = np.random.choice(['low', 'center', 'high'], p=[1/3, 1/3, 1/3])
        if flag_B == 'low': 
            B = np.random.randint(0, math.ceil((S) / 3))
        elif flag_B == 'center':
            B = np.random.randint(math.ceil((S) / 3), math.ceil((S) / 3) * 2)
        else:
            B = random.randint(math.ceil((S) / 3) * 2, (S))
        #print(flag_B, B)

        if A + B > S or B == 0:
            continue
        else:
            # Cの値を決める。これは総和150から引くだけ
            C = S - (A + B)
            i = i + 1
    data_list.append([A, B, C])

# パンダスデータフレーム形式へ変換
df = pd.DataFrame(data_list, columns=parameter_name_list)
# 各列の和を計算して、列名「sum」に格納する
df['sum'] = df['A'] + df['B'] + df['C']
# インデックス番号を1からに振り直す
df.index = np.arange(1, len(df) + 1)
df.to_csv('data1.csv', index=True)
df


# In[2]:


import matplotlib.pyplot as plt

# ヒストグラムで分布を確認する
fig = plt.figure(figsize=(6,6))
ax = fig.gca()
df.hist(ax = ax)
plt.tight_layout()

▼2. 各要素の生成する乱数の範囲をある程度指定する場合
下図のように、Aは正規分布に近いといった分布を恣意的に作成できます。当然ですが、和が一定のためにBとCはそれにならった分布になります。

A

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

parameter_name_list = ['A', 'B', 'C']
S = 150 # 制約条件 S＝A＋B＋Cが固定値
N = 2000 # 水準数

# 乱数のインスタンス生成
rng = np.random.default_rng(1) # ()内はランダムシード
rng


# In[2]:


# パラメータAの生成
parameter_A = rng.integers(15, 85, size=N) # 15以上85未満の整数の乱数N個
parameter_A


# In[3]:


# パラメータBの生成
parameter_B = rng.integers(20, 55, size=N) # 20以上55未満の整数の乱数N個
parameter_B


# In[4]:


# パラメータCの生成
parameter_C = rng.integers(15, 35, size=N) # 15以上35未満の整数の乱数N個
parameter_C


# In[5]:


# 上記で生成した各パラメータをひとつのパラメータセットに合成する
np_list = [parameter_A, parameter_B, parameter_C]
param_arr = np.vstack(np_list).T
param_arr


# In[6]:


# パラメータA,B,Cの合計がSになるように変換する。ここで、整数ではなくなってしまう
arr = param_arr * S / param_arr.sum(axis=1)[:, np.newaxis]
arr


# In[7]:


# Sが指定した値になっているかを確認する
arr.sum(axis=1)


# In[8]:


# 四捨五入して整数にする
arr2 = np.floor(arr)
arr2


# In[9]:


# pandasデータフレーム型変換する
df = pd.DataFrame(arr2, columns=parameter_name_list)
df


# In[10]:


df['sum'] = df.sum(axis=1)
df


# In[11]:


# 和がSになるように最後のパラメータ（ここではC）の帳尻を合わせる
df2 = df.copy()
df2[parameter_name_list[-1]] = df[parameter_name_list[-1]] + (S - df['sum'])
df3 = df2.drop(['sum'], axis=1)
df3['sum'] = df3.sum(axis=1)
df3


# In[12]:


# csvファイルに出力する
df3.to_csv('data2.csv', index=False)


# In[13]:


# ヒストグラムで分布を確認する
fig = plt.figure(figsize=(6,8))
ax = fig.gca()
df3.hist(ax = ax)
plt.tight_layout()

以上

<広告>

リンク