# 2018 Q2(4)  2024.10.8

import numpy as np
import matplotlib.pyplot as plt

# パラメータの設定
N_values = [50, 70, 100]  # 総球数 N のバリエーション
M_steps = 10  # 赤球のステップ数
n = 10  # 抽出回数
n_trials = 10000  # シミュレーションの試行回数

# グラフ描画の準備
plt.figure(figsize=(12, 8))

for N in N_values:
    M_values = np.arange(M_steps, N + 1, M_steps)  # 赤球数をステップで変化
    
    # 理論値の計算
    E_X_theory = n * M_values / N
    V_X_theory = n * M_values / N * (N - M_values) / N * (N - n) / (N - 1)
    
    # シミュレーションによる期待値と分散の計算
    E_X_simulation = []
    V_X_simulation = []
    
    for M in M_values:
        X_values = []
        for _ in range(n_trials):
            balls = [1] * M + [0] * (N - M)
            drawn_balls = np.random.choice(balls, size=n, replace=False)
            X_values.append(np.sum(drawn_balls))  # 赤球の数
        
        E_X_simulation.append(np.mean(X_values))
        V_X_simulation.append(np.var(X_values))
    
    # 期待値のグラフに追加描画
    plt.subplot(2, 1, 1)  # 期待値のグラフ (上側)
    plt.plot(M_values, E_X_theory, label=f'理論値 (N={N})', linestyle='--')
    plt.plot(M_values, E_X_simulation, label=f'シミュレーション (N={N})', marker='o', linestyle='None')
    
    # 分散のグラフに追加描画
    plt.subplot(2, 1, 2)  # 分散のグラフ (下側)
    plt.plot(M_values, V_X_theory, label=f'理論値 (N={N})', linestyle='--')
    plt.plot(M_values, V_X_simulation, label=f'シミュレーション (N={N})', marker='o', linestyle='None')

# グラフの設定
plt.subplot(2, 1, 1)
plt.xlabel('赤球の数 (M)')
plt.ylabel('期待値 E[X]')
plt.title('異なる N における期待値 E[X] の理論値とシミュレーション結果の比較')
plt.legend()
plt.grid(True)

plt.subplot(2, 1, 2)
plt.xlabel('赤球の数 (M)')
plt.ylabel('分散 V[X]')
plt.title('異なる N における分散 V[X] の理論値とシミュレーション結果の比較')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()