HHKB プログラミングコンテスト 2020 D - Squares

AtCoder

問題

D - Squares

解法

解説: Editorial - HHKB Programming Contest 2020

上記解説手順で解いてみる

#include <bits/stdc++.h>
using namespace std;
#define REP(i,n) for(int i=0;i<n;i++)
typedef long long ll;
const int mod = 1000000007;
// mintは省略 https://github.com/atcoder/ac-library

int main() {
  int t;
  cin >> t;
  REP(i,t) {
    ll n,a,b;
    cin >> n >> a >> b;

    // N正方形内に赤と青の正方形を置く全パターンを求める
    mint ax = n-a+1; // N内にaの区間を置くパターン
    mint bx = n-b+1; // N内にbの区間を置くパターン
    mint all = ax*ax*bx*bx;

    // 赤い正方形の内部と青い正方形の内部が重なるケースを求める
    // x4: 空白マスはn-a-b.　(空白)(赤)(空白)(青)(空白) となるため空白+(赤と青分の+2)から2つを選ぶ組み合わせを求める
    mint x4 = a+b <= n ? (n-a-b+2)*(n-a-b+1)/2 : 0;
    mint x3 = x4*2;
    mint x2 = ax*bx-x3;
    mint x1 = x2*x2;

    // 全体から内部が重なるケースを引く
    mint ans = all-x1;
    cout << ans << endl;
  }
}