アルゴリズムと数学 066

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

各カードに書かれた整数をa, b, cとして、aを固定します。aが1とNから遠いところなら場合の数は同じです。それを含めて総当たりにすればよいです。

// Three Cards
#![allow(non_snake_case)]

use std::cmp::min;


//////////////////// 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() -> (i64, i64) {
    let v = read_vec();
    let N = v[0];
    let K = v[1];
    (N, K)
}

fn is_valid_triplet(a: i64, b: i64, c:i64, K: i64) -> bool {
    (a - b).abs() >= K || (a - c).abs() >= K || (b - c).abs() >= K
}

fn count_all_fixed_a(a: i64, N: i64, K: i64) -> i64 {
    let mut counter: i64 = 0;
    for b in 1..N+1 {
        for c in 1..N+1 {
            if is_valid_triplet(a, b, c, K) {
                counter += 1
            }
        }
    }
    counter
}

fn count_all(N: i64, K: i64) -> i64 {
    (1..N+1).map(|a| count_all_fixed_a(a, N, K)).sum::<i64>()
}

fn f(N: i64, K: i64) -> i64 {
    if N < 2*K-1 {
        count_all(N, K)
    }
    else {
        let mut counter = 0;
        let center = N*N - (K*2-1).pow(2) + count_all_fixed_a(K, K*2-1, K);
        for a in 1..N+1 {
            if K <= a && a <= N - K + 1 {
                counter += center
            }
            else {
                let b = min(a, N-a+1);
                let r = K + b - 1;
                counter += N*N - r*r + count_all_fixed_a(b, r, K)
            }
        }
        counter
    }
}

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