パ研合宿コンペティション 3日目 F - ミックスジュース

問題文は閲覧注意です。
F - ミックスジュース

問題

「A から B までの整数から 1 個以上を選んだ時の和として得られる整数は何通りあるかmod 10^9+7で答えよ」という問いに Q 回答えよ。
1 <= A <= B <= 10^9
Q <= 10^5

解法

AからBまでの整数のうちi個使う場合、和はi(2A+i-1)/2からi(2B-i+1)/2までの整数のどれかである。
これから、i=1,...,B-A+1について閉区間[i(2A+i-1)/2,i(2B-i+1)/2]の和集合をとり、そこに含まれる整数の個数を数えれば良い。
当然左端も右端もiにしたがって単調増加するので、求める値は(B-A+1番目の右端) - (1番目の左端) + 1 - \sum_{i = 1}^{B - A} max(0, (i + 1番目の左端) - (i番目の右端) - 1)である。
第1項と第2項は簡単。第3項については、(i + 1番目の左端) - (i番目の右端) - 1がiの2次式であり、この2次式のi=0とi=B-Aでの値がA-1であることがわかるので、これが0以上である範囲を求めればよい。そのような範囲はi=(B-A)/2を軸として点対称である。

実装上の注意点

  • 端の扱いに気をつける
  • 2次方程式の解は二分探索で求める

提出: Submission #3878249 - パ研合宿コンペティション 3日目 (Rust)

// ModInt 省略
fn calc(a: i64, b: i64) -> ModInt {
    let one = ModInt::new(1);
    let two = ModInt::new(2);
    let am = ModInt::new(a);
    let bm = ModInt::new(b);
    let mut tot = (bm + am + one) * (bm - am) * two.inv() + one;
    let half = (b - a) / 2;
    let mut pass = 0;
    let mut fail = half + 1;
    while fail - pass > 1 {
        let mid = (fail + pass) / 2;
        if mid * (a - b + mid) + a - 1 >= 0 {
            pass = mid;
        } else {
            fail = mid;
        }
    }
    if pass == half {
        // i^2 + (a - b)i + a - 1 (i = 1 .. b - a)
        let bma = bm - am;
        tot -= bma * (bma + one) * (two * bma + one) * ModInt::new(6).inv();
        tot += bma * bma * (bma + one) * two.inv();
        tot -= (am - one) * bma;
    } else {
        // 2 * (i^2 + (a - b)i + a - 1) (i = 1 .. pass) + (a - 1)
        let x = ModInt::new(pass);
        tot -= x * (x + one) * (two * x + one) * ModInt::new(3).inv();
        tot -= (am - bm) * x * (x + one);
        tot -= two * (am - one) * x;
        tot -= am - one;
    }
    tot
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        ab: [(i64, i64)],
    }
    for (a, b) in ab {
        puts!("{}\n", calc(a, b));
    }
}

まとめ

この問題文で中高生が参加するコンテストに出題するのはまずくないですか…?