Python3で解く AtCoder Regular Contest 107 B - Quadruple

まず、 $a, b, c, d$ すべてのパターンを一気に考えようとせず、 $1 \leq a, b \leq n$ かつ $a + b = x$ を満たす $a, b$ の組み合わせ数を出せるようにする。その後条件を詰めていく。

結果的に、問題は $(a+b) - (c+d) = k$ なので、 $(a+b) = x$ かつ $(c+d) = x-k$ として考えることができる。

概要

問題

整数 $n, k$ が与えられる。
4つの整数の組 $(a, b, c, d)$ であって，以下の条件を両方満たすものは何個あるか？
$1 \leq a, b, c, d \leq n$
$a + b - c - d = k$

数の大きさの制約は問題を見てもらう。

解くときに考えた内容

問題を $1 \leq a, b, c, d \leq n$ かつ $(a+b) - (c+d) = k$ なので、 $1 \leq a, b, c, d \leq n$ かつ $(a+b) = x$ かつ $(c+d) = x-k$ とまず捉える。
とすると、以下の条件を満たす $f(n, x)$ を解ければ良いとわかる。
- $1 \leq a, b, c, d \leq n$ かつ $2 \leq x \leq 2n$
例えば $n = 6, k = 3$ のとき
- $x=5$ における $a, b$ の組み合わせ数 $\times\ x-k = 2$ における $c, d$ の組み合わせ数（ $4 \times 1$ ）
- $x=6$ における $a, b$ の組み合わせ数 $\times\ x-k = 3$ における $c, d$ の組み合わせ数（ $5 \times 2$ ）
- $x=7$ における $a, b$ の組み合わせ数 $\times\ x-k = 4$ における $c, d$ の組み合わせ数（ $6 \times 3$ ）
- $x=8$ における $a, b$ の組み合わせ数 $\times\ x-k = 5$ における $c, d$ の組み合わせ数（ $5 \times 4$ ）
- $x=9$ における $a, b$ の組み合わせ数 $\times\ x-k = 6$ における $c, d$ の組み合わせ数（ $4 \times 5$ ）
- $x=10$ における $a, b$ の組み合わせ数 $\times\ x-k = 7$ における $c, d$ の組み合わせ数（ $3 \times 6$ ）
- $x=11$ における $a, b$ の組み合わせ数 $\times\ x-k = 8$ における $c, d$ の組み合わせ数（ $2 \times 5$ ）
- $x=12$ における $a, b$ の組み合わせ数 $\times\ x-k = 9$ における $c, d$ の組み合わせ数（ $1 \times 4$ ）
  ここまで考えればよい。
肝心の、 $f(n, x)$ は計算式では $min(k-1, 2*n - (k-1))$ 以下のように表せる。
- $n$ の制約がなければ、 $k-1$ だが、 $n$ の制約下だと、 $2*n - (k-1)$ のどちらか小さい方になる。
- また、うまく条件分岐すれば良いが、面倒なので、 $min(k-1, 2*n - (k-1))$ の値が $0$ 以下未満になったら $0$ となるようにする。

コード

def a_b_x(n, k):
    # a+b=xかつ1<=a, b<=nを満たす個数
    # 数えるべき部分をif文で書くべきだが、面倒なので範囲外は0とする。
    ans = max(0, min(k-1, 2*n - (k-1)))
    return ans            
n, k = map(int, input().split())
cnt = 0
for i in range(2, 2*n+1):
    cnt += a_b_x(n, i) * a_b_x(n, i-k)
print(cnt)

参考になる書籍

問題解決力を鍛える!アルゴリズムとデータ構造

けんちょんさんこと、大槻さんの書籍である、「問題解決力を鍛える!アルゴリズムとデータ構造」ご本人の記事の中で難易度は、螺旋本ということで初学者でも非常に結構読みやすい!!!（実際に、けんちょんさんの記事はC++に疎い自分でも、その思考プロセスを学べるくらい丁寧に書かれているし、この書籍も非常に読みやすい）どのようなステータスの人にもおすすめの一冊になっている。