アルゴリズムと数学 091(2)

https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_br

しらみつぶしでも十分速いのですが、O(1)でも計算可能です。
 1 \le a, b, c
 a + b + c = X
の条件だと、正三角形の領域になります。ここでN以下という条件をつけると、NがX-2より小さければ、正三角形の頂点付近の正三角形が落とされて、六角形の領域になります。さらにNが小さければ、切り落とされすぎて再び正三角形の領域になります。例えば、N=4, X=10の場合です。そこから、a, b, cのどれかが同じときを除いて、6で割ります。

// How Many Ways?
#![allow(non_snake_case)]

use std::cmp::{min, max};


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (i32, i32) {
    let v = read_vec();
    let N = v[0];
    let X = v[1];
    (N, X)
}

fn T(N: i32) -> i32 {
    N * (N + 1) / 2
}

fn f(N: i32, X: i32) -> i32 {
    if X < 6 {
        0
    }
    else if 3*N + 3 < X {
        0
    }
    else if 2*N+1 > X {
        let T1 = T(X-2);                        // 大きな三角形
        let T2 = T((X-N-2).max(0));             // 端の三角形
        let D1 = (X-1)/2-max(2, X-N+1)/2+1;     // 2つ同じ
        let D2 = if X%3 == 0 { 2 } else { 0 };  // 3つ同じ
        (T1 - 3*T2 - 3*D1 + D2) / 6
    }
    else {
        f(3*N-X+1, 3*N-X+3)
    }
}

fn main() {
    let (N, X) = read_input();
    println!("{}", f(N, X))
}