[Python] ABC003 D

D – AtCoder社の冬

import math

# 順列
def P(n, r):
    return math.factorial(n)//math.factorial(n-r)
# 組み合わせ
def C(n, r):
    return P(n, r)//math.factorial(r)
# x*yマスの中にデスクとラックと空白を置くパターンが何通りあるかを返す
def allocation_d_l(x, y):
    #包除原理を使う際にあり得ない値を除去しておく
    if x<0 or y<0:
        return 0
    elif x*y < d+l:
        return 0
    else:
        return C(x * y, d) * C(x * y - d, l)

# AtCoder社の社員室の区画の行数r、列数c
r, c = map(int, input().split())
# 社員室の壁に囲まれた部分の区画の行数x、列数y
x, y = map(int, input().split())
# d デスクの数 l サーバーラックの数
d, l = map(int, input().split())

# r*cマスの中に、x*yマスを配置するパターン
pattern_wall = (r - x + 1) * (c - y + 1)

# x*yマスの中がデスクとラックで埋まる場合 100点
if x*y == d+l:
    pattern_d_l = C(x*y, d)
# xyマスの中に空白が入る場合 101点
elif x*y > d+l:
    # xyマスの中にデスクとラックと空白を配置する全パターン数
    pattern_d_l = allocation_d_l(x, y)
    # 全パターンから取り除くパターン、つまりいずれかの辺が空白であるパターンを、包除原理で求める。
    pattern_eliminate = allocation_d_l(x-1, y)*2 + allocation_d_l(x, y-1)*2
    pattern_eliminate -= allocation_d_l(x-1, y-1)*4 + allocation_d_l(x-2, y) + allocation_d_l(x, y-2)
    pattern_eliminate += allocation_d_l(x-1, y-2) * 2 + allocation_d_l(x-2, y-1) * 2 
    pattern_eliminate -= allocation_d_l(x-2, y-2)
    # 全パターンから取り除くパターンを引く
    pattern_d_l = pattern_d_l - pattern_eliminate

divisor = 10**9+7
print(pattern_wall * pattern_d_l % divisor)

100点は取れたけど、101点が取れなかった…。

メモ

  • 包除原理を使う。
  • 包除原理 解ける数え上げの範囲を広げよう とても分かりやすい資料!
  • 考え方としては、デスクかラックが必ず上下左右に1つはなければならないので、つまり、(全パターン)から((上が全て空白、以下Aとする) ∨(左が全て空白、以下Bとする) ∨(下が全て空白、以下Cとする) ∨(右が全て空白、以下Dとする))を引けばよい。
  • (A∨B∨C∨D) は、包除原理より、 (A+B+C+D) – ((A∧B)+(A∧C)+(A∧D)+(B∧C)+(B∧D)+(C∧D) ) + ((A∧B∧C)+(A∧B∧D)+(A∧C∧D)+(B∧C∧D)) – (A∧B∧C∧D)と表すことができる。
  • allocationとして定義した関数で組み合わせを求める場合、エラーを防ぐため、事前にあり得ない値を除去しておく必要がある。