[Python] 이항분포 그래프 그리기

import sys
from functools import lru_cache
import math
import numpy as np
import matplotlib.pyplot as plt


@lru_cache(None)
def ncr(n, r):
    """조합. 재귀식을 이용함.
    n이 커지면 스택오버플로우 발생.
    """

    if r in (0, n):
        return 1
    return ncr(n - 1, r) + ncr(n - 1, r - 1)


def bidist(n, p):
    """이항분포"""

    q = 1 - p
    dist = np.array([ncr(n, k) * (p ** k) * (q ** (n - k)) for k in range(n + 1)])
    return dist


def bicoeff(n):
    """이항계수. 파스칼의 삼각형."""

    l = [1]
    for _ in range(n):
        l = [a + b for a, b in zip(l + [0], [0] + l)]
    return l


def bidist2(n, p):
    """이항분포"""

    q = 1 - p
    coeff = bicoeff(n)
    dist = np.array([coeff[k] * (p ** k) * (q ** (n - k)) for k in range(n + 1)])
    return dist


def show_plot(n=120, p=0.38):
    plt.grid(True)
    plt.plot(bidist2(n, p))
    m, s = n * p, math.sqrt(n * p * (1 - p))
    plt.axvline(x=m, color="r", alpha=0.6)
    plt.axvline(x=(m - s), color="grey", alpha=0.6)
    plt.axvline(x=(m + s), color="grey", alpha=0.6)
    plt.title("Binomial Distribution $ B(%d, %lf) $" % (n, p))
    plt.show()


if __name__ == "__main__":
    if len(sys.argv) == 3:
        ## 실행인자를 주어서, n, p 값을 바꾸어가며 분포를 볼 수 있다.
        n, p = int(sys.argv[1]), float(sys.argv[2])
        if n > 1 and 0 < p < 1:
            show_plot(n, p)
            sys.exit()
    ## 실행인자가 없거나 조건에 맞지 않으면, 
    ## n=120, p=0.38 기본값의 분포를 보여줌.
    show_plot()

 

n=120, p=0.38 일 때의 분포

 

n=7, p=0.38 일 때의 분포