系

N個の飴を棚に配る
棚には3種
- 必須棚：1つ以上の飴配置が必要。M個あるとする
- ペア棚：2つ1組。ペア1つのうちどちらか1つに1個以上の飴配置が必要。P組あるとする（=棚は2*P個）
- 許容棚：飴を配っても配らなくても良い自由な棚。A個あるとする
全ての飴・棚にはラベルがあり区別可能
N、M、P、Aが与えられたときの場合の数を計算する

制約の解釈

包除原理を用いる。

必須棚：「1つも配られない」数を全数から引く
ペア棚：「どちらにも1つもない」数を全数から引く

ベン図、と言った通り、条件が複数あって、それぞれが排反ではない（条件1と条件2の同時存在があり得る）場合に包除原理を使う。
重なり部分のかぞえすぎをキャンセルする。

包除原理を使用せずにバラバラに数えてみた
先に必要分最小限の配置組合せのみを考慮し、残りの飴の数で$(M+A+2P)^(N-M-P)$で自由は以下を計算、
数え過ぎを引くことを検討。
数え過ぎるケースは、最小数を配置した場合に対し、自由配置の飴が再度割り振られる場合で且つ自由配置分が最小数配置の飴ラベルよりも小さい値の時に数え過ぎとなる。
Mに2個の飴が配置さえっる場合の数は、階段状（三角形の積み重ね）的イメージで計算できたが、
Mに3個以上の配下がなされる状況が図でイメージしにくく、Pに関する検討も難しくて断念した

包除原理

2つの集合の時：

|A\cup B| = |A| + |B| - |A \cap B|

3つの集合の時、今度は引きすぎるので加える操作が入る。

∣A\cup B\cup C∣=∣A∣+∣B∣+∣C∣−∣A\cap B∣−∣A \cap C∣−∣B \cap C∣+∣A\cap B\cap C∣

同様にして一般のn個の集合の時は以下のように表す。

\begin{align}
\left| \bigcup_{i=1}^n A_i \right|
&= \sum_{1 \le i \le n} |A_i|
\; - \sum_{1 \le i < j \le n} |A_i \cap A_j|
\; + \sum_{1 \le i < j < k \le n} |A_i \cap A_j \cap A_k|
\; - \cdots
\; + (-1)^{n+1} \left| \bigcap_{i=1}^n A_i \right|
\\
\\
\left| \bigcup_{i=1}^n A_i \right|
&= \sum_{k=1}^{n} (-1)^{k+1}
\sum_{1 \le i_1 < i_2 < \cdots < i_k \le n}
\left| A_{i_1} \cap A_{i_2} \cap \cdots \cap A_{i_k} \right|
\end{align}

最初にすべての集合の要素数を足す
2つずつの交差部分を引く
3つずつの交差部分を足す
交互に繰り返し$(-1)^i$

コードで計算したい。

計算

理論式：

使用する棚、使用しない棚についてカウントしていく
これをコードに落とし込む。
一方、検算が大変なので、愚直に数え上げるコードで検算するようにする。

数式の展開：

\sum_{S\subseteq M ,T \subseteq P} \dots

$\sum$の部分はfor文の2重ループで書く。

\begin{align}
&(M+2P+A - |S| - 2|T|)^N  
\\
& |X| := 集合Xの要素数
\end{align}

|S|、|T|は「使用しない」必須棚、ペア棚の数（制約違反、ペナルティに関連）
forループで足し算をしつつ、使用する棚数で配荷数を数えていく、という事を行う。
包除原理で、数え過ぎ（ベン図の重なり部分）を足し引きして数を整える

from itertools import product
# product : 直積。集合同士の要素を取り出して、考え得るすべての組み合わせを得る操作

def theoretical_distribution_count(N, M, P, A):
    """
    定式化して、理論式から計算
    Args:
      N:配荷するアイテム数
      M：必須棚数
      P：ペア棚数
      A：許容棚数
    """

    total = 0

    # 必須棚の空の組み合わせ
    for s in range(1 << M):
        num_empty_m = bin(s).count('1')

        # ペア棚の両方空のペアの組み合わせ
        for t in range(1 << P):  
            num_empty_p = bin(t).count('1')
            usable_shelves = M + 2 * P + A - num_empty_m - 2 * num_empty_p
            if usable_shelves < 0:
                continue
            total += (-1) ** (num_empty_m + num_empty_p) * (usable_shelves ** N)
    return total

# 数え上げる
def count_valid_distributions(N:int, M:int, P:int, A:int):
    """
    実際にカウントをして求める
    Args:
      N:配荷するアイテム数
      M：必須棚数
      P：ペア棚数
      A：許容棚数
    """

    total_valid = 0

    # 全棚ラベル：必須棚 m0〜、ペア棚 p0a, p0b, p1a, ..., 許容棚 a0〜
    shelves = []

    # 棚インデックス
    must_indices = list(range(M))  #=[0,1,2,3,,,]
    pair_indices = []

    for i in range(P):
      # Mから数え始めて2個ずつ増やす
      pair_indices.append((M + 2*i, M + 2*i + 1))

    # pな軸、M+2Pから始めてAの数だけ拡張したリスト
    allow_indices = list(range(M + 2*P, M + 2*P + A))

    shelves = must_indices + [i for pair in pair_indices for i in pair] + allow_indices
    num_shelves = len(shelves)

    # 全ての飴に対する棚の割り当てを列挙（N個の飴に対して num_shelves 通り）
    for assignment in product(range(num_shelves), repeat=N):
        # 棚ごとのカウント
        count = [0] * num_shelves
        for s in assignment:
            count[s] += 1

        # 条件チェック
        valid = True

        # ① 必須棚すべてに1個以上
        for i in must_indices:
            if count[i] == 0:
                valid = False
                break

        if not valid:
            continue

        # ② ペア棚：各ペアで「少なくとも片方に1個以上」
        for i, j in pair_indices:
            if count[i] == 0 and count[j] == 0:
                valid = False
                break

        if valid:
            total_valid += 1

    return total_valid


# テストケース
N=4
M=2
P=2
A=10

result_count_up = count_valid_distributions(N, M, P, A)
result_theoretical = theoretical_distribution_count(N, M, P, A)
print(result_count_up)
print(result_theoretical)