「○通り求めよ」系問題で初手に全探索を行うことを止めた (AGC015 A - A+...+B Problem)

初めてAGCのA問題の茶色級*1を解くことができたので、解法を記載する。

問題

atcoder.jp

問題は「最小値A、最大値Bの長さNの数列が与えられる。この数列の総和は何通りあるか？」である。
制約はN, A, Bともに1〜10^9をとる。

「○通り求めよ」系の問題は一瞬身構えてしまう。
初心者にありがちな「何でも全探索で解く」をやりがちな問題の聞き方であり、またそんな解き方をすると即座にTLEを叩きつけられる問題でもある。

考え方

当然全探索を行うと計算量は天文学的数値を取るため、この考え方は捨てる。
この問題は初めに和を考えるよりも、数列の0番目とN-1番目を除く数字の並べ方を考えるほうが良い。

たとえばN = 4, A = 4, B = 6の場合、以下の図のような○に4, 5, 6のいずれかを入れることを考える。

この場合、それぞれの○に4〜6の3つの数値を入れる場合の数を計算すると、和がかぶるパターンが出る。

そのため、以下のように考える。
○○に入る数値は最小値「44」、最大値「66」である。和がかぶらないようにするためには以下の操作を行う。
操作: 右の○から順番に数値をインクリメントする
数値が最大値となるまで上記の操作を繰り返す。

この操作を行った回数が、今回求める値である。

この操作回数は「◯に入りうる数値の個数(B - A + 1)」「◯の数(N - 2)」を考えると、以下のような式で計算することができる。
(B - A + 1)(N - 2) - ((N - 2) - 1)
= (B - A + 1)(N - 2) - (N - 2) + 1
= (B - A)(N - 2) + 1

あとは、提示されている例にあるコーナーケースに気をつけながら実装する。

提出コード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

int main(){
    ll N, A, B; cin >> N >> A >> B;

    if(A > B) {
        cout << 0 << endl;
    } else if(N == 1 && B - A != 0) {
        cout << 0 << endl;
    } else if(N == 1 && B - A == 0) {
        cout << 1 << endl;
    } else {
        cout << (B - A)(N - 2) + 1 << endl;
    }
}

*1:AtCoder Problems参照