#497889 (Rust) No.800 四平方定理

提出ソース

結果

問題	No.800 四平方定理
ユーザー	yoshig0731
提出日時	2020-06-17 19:45:49
言語	Rust (1.77.0)
結果	AC
実行時間	686 ms / 2,000 ms
コード長	20,063 bytes
コンパイル時間	1,407 ms
コンパイル使用メモリ	180,552 KB
実行使用メモリ	142,408 KB
最終ジャッジ日時	2023-09-16 11:39:14
合計ジャッジ時間	11,727 ms
ジャッジサーバーID （参考情報）	judge14 / judge13

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	5 ms 4,380 KB
testcase_01	AC	8 ms 4,376 KB
testcase_02	AC	7 ms 4,380 KB
testcase_03	AC	7 ms 4,380 KB
testcase_04	AC	6 ms 4,376 KB
testcase_05	AC	7 ms 4,380 KB
testcase_06	AC	10 ms 4,380 KB
testcase_07	AC	7 ms 4,376 KB
testcase_08	AC	9 ms 4,376 KB
testcase_09	AC	8 ms 4,376 KB
testcase_10	AC	310 ms 71,920 KB
testcase_11	AC	312 ms 70,848 KB
testcase_12	AC	317 ms 70,856 KB
testcase_13	AC	318 ms 70,860 KB
testcase_14	AC	331 ms 70,604 KB
testcase_15	AC	319 ms 70,468 KB
testcase_16	AC	315 ms 71,120 KB
testcase_17	AC	325 ms 71,124 KB
testcase_18	AC	317 ms 70,852 KB
testcase_19	AC	313 ms 70,580 KB
testcase_20	AC	1 ms 4,376 KB
testcase_21	AC	1 ms 4,376 KB
testcase_22	AC	361 ms 70,132 KB
testcase_23	AC	686 ms 141,044 KB
testcase_24	AC	621 ms 140,192 KB
testcase_25	AC	670 ms 139,404 KB
testcase_26	AC	1 ms 4,376 KB
testcase_27	AC	1 ms 4,376 KB
testcase_28	AC	585 ms 142,408 KB
testcase_29	AC	625 ms 140,512 KB
testcase_30	AC	672 ms 139,932 KB
testcase_31	AC	585 ms 139,640 KB
testcase_32	AC	675 ms 141,824 KB

権限があれば一括ダウンロードができます

コンパイルメッセージ

warning: unused import: `std::f64::consts::PI`
 --> Main.rs:3:5
  |
3 | use std::f64::consts::PI;
  |     ^^^^^^^^^^^^^^^^^^^^
  |
  = note: `#[warn(unused_imports)]` on by default

warning: unused import: `std::str::*`
 --> Main.rs:6:5
  |
6 | use std::str::*;
  |     ^^^^^^^^^^^

warning: unused macro definition: `min`
  --> Main.rs:63:14
   |
63 | macro_rules! min {
   |              ^^^
   |
   = note: `#[warn(unused_macros)]` on by default

warning: unused macro definition: `chmin`
  --> Main.rs:75:14
   |
75 | macro_rules! chmin {
   |              ^^^^^

warning: unused macro definition: `max`
  --> Main.rs:87:14
   |
87 | macro_rules! max {
   |              ^^^

warning: unused macro definition: `chmax`
  --> Main.rs:99:14
   |
99 | macro_rules! chmax {
   |              ^^^^^

warning: the item `Write` is imported redundantly
   --> Main.rs:134:13
    |
4   | use std::io::*;
    |     ---------- the item `Write` is already imported here
...
134 |         use std::io::Write;
    |             ^^^^^^^^^^^^^^

warning: the item `Read` is imported redundantly
   --> Main.rs:139:13
    |
4   | use std::io::*;
    |     ---------- the item `Read` is already imported here
...
139 |         use std::io::Read;
    |             ^^^^^^^^^^^^^

warning: the item `BinaryHeap` is imported redundantly
   --> Main.rs:265:9
    |
2   | use std::collections::*;
    |     ------------------- the item `BinaryHeap` is already imported here
...
265 |     use std::collections::BinaryHeap;
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^

warning: unused variable: `tar`
   --> Main.rs:518:19
    |
518 |     fn act(&self, tar: &Self::T, x: usize) -> Self::T {
    |                   ^^^ help: if this is intentional, prefix it with an underscore: `_tar`
    |
    = note: `#[warn(unused_variables)]` on by default

warning: unused variable: `x`
   --> Main.rs:518:34
    |
518 |     fn act(&self, tar: &Self::T, x: usize) -> Self::T {
    |                                  ^ help: if this i

ソースコード

raw source code

use std::cmp::*;
use std::collections::*;
use std::f64::consts::PI;
use std::io::*;
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};
use std::str::*;
use std::string::*;

macro_rules! read {
    (($($t:tt),*)) => {
        ( $(read!($t)),* )
    };
    ([[$t:tt; $len1:expr]; $len2:expr]) => {
        (0..$len2).map(|_| read!([$t; $len1])).collect::<Vec<_>>()
    };

    ([$t:tt; $len:expr]) => {
        (0..$len).map(|_| read!($t)).collect::<Vec<_>>()
    };

    (chars) => {
        read!(String).chars().collect::<Vec<char>>()
    };

    (usize1) => {
        read!(usize) - 1
    };

    ($t:ty) => {{
        let stdin = stdin();
        let stdin = stdin.lock();
        let token: String = stdin
            .bytes()
            .map(|c| c.unwrap() as char)
            .skip_while(|c| c.is_whitespace())
            .take_while(|c| !c.is_whitespace())
            .collect();

        token.parse::<$t>().unwrap()
    }};
}

macro_rules! input {
    (mut $name:ident: $t:tt, $($r:tt)*) => {
        let mut $name = read!($t);
        input!($($r)*);
    };

    (mut $name:ident: $t:tt) => {
        let mut $name = read!($t);
    };

    ($name:ident: $t:tt, $($r:tt)*) => {
        let $name = read!($t);
        input!($($r)*);
    };

    ($name:ident: $t:tt) => {
        let $name = read!($t);
    };
}

macro_rules! min {
    ($a:expr $(,)*) => {{
        $a
    }};
    ($a:expr, $b:expr $(,)*) => {{
        std::cmp::min($a, $b)
    }};
    ($a:expr, $($rest:expr),+ $(,)*) => {{
        std::cmp::min($a, min!($($rest),+))
    }};
}

macro_rules! chmin {
    ($base:expr, $($cmps:expr),+ $(,)*) => {{
        let cmp_min = min!($($cmps),+);
        if $base > cmp_min {
            $base = cmp_min;
            true
        } else {
            false
        }
    }};
}

macro_rules! max {
    ($a:expr $(,)*) => {{
        $a
    }};
    ($a:expr, $b:expr $(,)*) => {{
        std::cmp::max($a, $b)
    }};
    ($a:expr, $($rest:expr),+ $(,)*) => {{
        std::cmp::max($a, max!($($rest),+))
    }};
}

macro_rules! chmax {
    ($base:expr, $($cmps:expr),+ $(,)*) => {{
        let cmp_max = max!($($cmps),+);
        if $base < cmp_max {
            $base = cmp_max;
            true
        } else {
            false
        }
    }};
}

#[derive(Eq, PartialEq, Clone, Debug)]
pub struct Rev<T>(pub T);

impl<T: PartialOrd> PartialOrd for Rev<T> {
    fn partial_cmp(&self, other: &Rev<T>) -> Option<std::cmp::Ordering> {
        other.0.partial_cmp(&self.0)
    }
}

impl<T: Ord> Ord for Rev<T> {
    fn cmp(&self, other: &Rev<T>) -> std::cmp::Ordering {
        other.0.cmp(&self.0)
    }
}

struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> Self {
        IO(r, std::io::BufWriter::new(w))
    }

    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write(s.to_string().as_bytes()).unwrap();
    }

    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let s: String = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap() as char)
            .skip_while(|c| c.is_whitespace())
            .take_while(|c| !c.is_whitespace())
            .collect();

        return s.parse::<T>().ok().unwrap();
    }

    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }

    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

fn prime_factorization(n: i64) -> HashMap<i64, i64> {
    let mut n = n;
    let mut res = HashMap::new();
    let mut i = 2;
    while i * i <= n {
        while n % i == 0 {
            n /= i;
            let count = res.entry(i).or_insert(0);
            *count += 1;
        }
        i += 1;
    }

    if n > 1 {
        res.insert(n, 1);
    }

    return res;
}

//
// struct Combination {
//     MOD: i64,
//     fac: Vec<i64>,
//     fac_inv: Vec<i64>,
// }
//
// impl Combination {
//     pub fn new(n: i64) -> Self {
//         let MOD: i64 = 1_000_000_007;
//         let mut fac: Vec<i64> = vec![0; n as usize + 1];
//         let mut fac_inv: Vec<i64> = vec![0; n as usize + 1];
//
//         let get_inverse = |mut n: i64| -> i64 {
//             let (mut res, mut p) = (1, MOD - 2);
//
//             while p != 0 {
//                 if p & 1 == 1 {
//                     res = (res * n) % MOD;
//                 }
//                 n = (n * n) % MOD;
//                 p >>= 1;
//             }
//
//             return res;
//         };
//
//         fac[0] = 1;
//
//         for i in 1..n + 1 {
//             fac[i as usize] = (fac[i as usize - 1] * i) % MOD;
//         }
//
//         for i in 0..n + 1 {
//             fac_inv[i as usize] = get_inverse(fac[i as usize]);
//         }
//
//         Combination {
//             MOD: MOD,
//             fac: fac,
//             fac_inv: fac_inv,
//         }
//     }
//
//     fn nCr(&self, n: i64, r: i64) -> i64 {
//         if n < r {
//             return 0;
//         }
//
//         let a: i64 = self.fac[n as usize];
//         let b: i64 = self.fac_inv[(n - r) as usize];
//         let c: i64 = self.fac_inv[r as usize];
//         let bc: i64 = (b * c) % self.MOD;
//
//         return (a * bc) % self.MOD;
//     }
//
//     fn nPr(&self, n: i64, r: i64) -> i64 {
//         if n < r {
//             return 0;
//         }
//
//         let a: i64 = self.fac[n as usize];
//         let b: i64 = self.fac_inv[(n - r) as usize];
//
//         return (a * b) % self.MOD;
//     }
//
//     fn nHr(&self, n: i64, r: i64) -> i64 {
//         if n == 0 && r == 0 {
//             return 1;
//         }
//
//         return self.nCr(n + r - 1, r);
//     }
// }

#[derive(Copy, Clone, Eq, PartialEq)]
struct Edge {
    to: usize,
    cost: i64,
}

fn dijkstra(graph: &Vec<Vec<Edge>>, s: &usize) -> Vec<i64> {
    use std::collections::BinaryHeap;
    let mut dist = vec![1e18 as i64; graph.len()];
    let mut heap = BinaryHeap::new();
    dist[*s] = 0;
    heap.push(Rev((0, *s)));
    while let Some(Rev((cost, v))) = heap.pop() {
        if dist[v] < cost {
            continue;
        }

        for e in &graph[v] {
            if dist[e.to] <= dist[v] + e.cost {
                continue;
            }
            dist[e.to] = dist[v] + e.cost;
            heap.push(Rev((dist[e.to], e.to)));
        }
    }
    return dist;
}

struct LCA {
    par: Vec<Vec<Option<usize>>>,
    dist: Vec<i64>,
}

impl LCA {
    pub fn new(graph: &Vec<Vec<usize>>, root: usize) -> LCA {
        let V = graph.len();
        let mut K = 1;

        while (1 << K) < V {
            K += 1;
        }

        let mut par = vec![vec![None; V]; K];
        let mut dist = vec![-1; V];
        let graph = graph.to_vec();

        fn dfs(
            v: usize,
            p: Option<usize>,
            d: i64,
            graph: &Vec<Vec<usize>>,
            par: &mut Vec<Vec<Option<usize>>>,
            dist: &mut Vec<i64>,
        ) {
            par[0][v] = p;
            dist[v] = d;
            for &to in &graph[v] {
                match p {
                    Some(p) => {
                        if to != p {
                            dfs(to, Some(v), d + 1, graph, par, dist)
                        }
                    }
                    None => dfs(to, Some(v), d + 1, graph, par, dist),
                }
            }
        }

        dfs(root, None, 0, &graph, &mut par, &mut dist);

        for k in 0..K - 1 {
            for v in 0..V {
                match par[k][v] {
                    Some(x) => par[k + 1][v] = par[k][x],
                    None => (),
                }
            }
        }

        LCA {
            par: par,
            dist: dist,
        }
    }

    pub fn query(&self, u: usize, v: usize) -> usize {
        let mut u = u;
        let mut v = v;
        if self.dist[u] < self.dist[v] {
            return self.query(v, u);
        }

        let K = self.par.len();
        for k in 0..K {
            if ((self.dist[u] - self.dist[v]) >> k & 1) == 1 {
                u = self.par[k][u].unwrap();
            }
        }

        if u == v {
            return u;
        }

        for k in (0..K - 1).rev() {
            if self.par[k][u] != self.par[k][v] {
                u = self.par[k][u].unwrap();
                v = self.par[k][v].unwrap();
            }
        }

        return self.par[0][u].unwrap();
    }
}

struct Doubling {
    N: usize,
    LOG: usize,
    next: Vec<Vec<Option<usize>>>,
}

impl Doubling {
    pub fn new(vec: &Vec<usize>, lim: &usize) -> Doubling {
        let N = vec.len();
        let lim = *lim;
        let mut LOG = 1;
        while (1 << LOG) < 2 * lim {
            LOG += 1;
        }

        let mut next = vec![vec![None; N]; LOG];

        for i in 0..N {
            next[0][i] = Some(vec[i]);
        }

        for k in 0..LOG - 1 {
            for i in 0..N {
                match next[k][i] {
                    Some(x) => next[k + 1][i] = next[k][x],
                    None => (),
                }
            }
        }

        Doubling {
            N: N,
            LOG: LOG,
            next: next,
        }
    }

    pub fn query(&self, i: usize, n: usize) -> usize {
        let mut i = i;
        for k in (0..self.LOG).rev() {
            if (n >> k) & 1 == 1 {
                i = self.next[k][i].unwrap();
            }
        }
        return i;
    }
}

struct UnionFind {
    par: Vec<i32>,
}

impl UnionFind {
    pub fn new(n: usize) -> UnionFind {
        UnionFind { par: vec![-1; n] }
    }

    pub fn root(&mut self, x: usize) -> usize {
        if self.par[x] < 0 {
            return x;
        } else {
            let a = self.par[x] as usize;
            self.par[x] = self.root(a) as i32;
            return self.par[x] as usize;
        }
    }

    pub fn same(&mut self, x: usize, y: usize) -> bool {
        self.root(x) == self.root(y)
    }

    pub fn merge(&mut self, x: usize, y: usize) -> bool {
        let mut x = self.root(x);
        let mut y = self.root(y);
        if x == y {
            return false;
        } else {
            if self.par[x] > self.par[y] {
                std::mem::swap(&mut x, &mut y);
            }
            self.par[x] += self.par[y];
            self.par[y] = x as i32;
            return true;
        }
    }

    pub fn size(&mut self, x: usize) -> i64 {
        let a = self.root(x);
        return -self.par[a] as i64;
    }
}

trait Monoid {
    fn id() -> Self;
    fn op(&self, rhs: &Self) -> Self;
}

#[derive(Clone, Debug)]
struct Min(i64);

#[derive(Clone, Debug)]
struct RangeUpdate(i64);

#[derive(Clone, Debug)]
struct RangeAdd(i64);

#[derive(Clone, Debug)]
struct Sum(i64);

impl Monoid for Sum {
    fn id() -> Self {
        Sum(0)
    }

    fn op(&self, rhs: &Self) -> Self {
        Sum(self.0 + rhs.0)
    }
}

impl Monoid for Min {
    fn id() -> Self {
        Min(std::i32::MAX as i64)
    }

    fn op(&self, rhs: &Self) -> Self {
        Min(std::cmp::min(self.0, (*rhs).0))
    }
}

trait Operator {
    type T;
    fn id() -> Self;
    fn op(&self, rhs: &Self) -> Self;
    fn act(&self, tar: &Self::T, x: usize) -> Self::T;
}

impl Operator for RangeUpdate {
    type T = Min;
    fn id() -> Self {
        RangeUpdate(0)
    }

    fn op(&self, rhs: &Self) -> Self {
        (*rhs).clone()
    }

    fn act(&self, tar: &Self::T, x: usize) -> Self::T {
        Min(self.0)
    }
}

impl Operator for RangeAdd {
    type T = Sum;
    fn id() -> Self {
        RangeAdd(0)
    }

    fn op(&self, rhs: &Self) -> Self {
        RangeAdd(self.0 + rhs.0)
    }

    // Sum(tar.0 + self.0 * x as i64)
    fn act(&self, tar: &Self::T, x: usize) -> Self::T {
        Sum(tar.0 + self.0 * x as i64)
    }
}

struct LazySegTree<M: Monoid, O: Operator<T = M>> {
    dat: Vec<M>,
    lzy: Vec<O>,
    need_update: Vec<bool>,
    size: usize,
}

impl<M: Monoid + Clone, O: Operator<T = M> + Clone> LazySegTree<M, O> {
    pub fn new(vec: Vec<M>) -> LazySegTree<M, O> {
        let n = vec.len();
        let mut sz = 1;
        while sz < n {
            sz <<= 1;
        }
        let mut dat = vec![M::id(); sz << 1];
        for i in 0..n {
            dat[i + sz] = vec[i].clone();
        }
        for k in (1..sz).rev() {
            dat[k] = dat[k << 1].op(&dat[(k << 1) + 1]);
        }

        return LazySegTree {
            dat: dat,
            lzy: vec![O::id(); sz << 1],
            need_update: vec![false; sz << 1],
            size: sz,
        };
    }

    pub fn eval(&mut self, l: usize, r: usize, k: usize) {
        if !self.need_update[k] {
            return;
        }
        self.dat[k] = self.lzy[k].act(&self.dat[k], r - l);
        if (k << 1) < self.dat.len() {
            self.lzy[k << 1] = self.lzy[k << 1].op(&self.lzy[k]);
            self.lzy[(k << 1) + 1] = self.lzy[(k << 1) + 1].op(&self.lzy[k]);
            self.need_update[k << 1] = true;
            self.need_update[(k << 1) + 1] = true;
        }
        self.lzy[k] = O::id();
        self.need_update[k] = false;
    }

    pub fn update(&mut self, a: usize, b: usize, x: O) {
        let sz = self.size;
        self._update(a, b, &x, 0, sz, 1);
    }

    pub fn _update(&mut self, a: usize, b: usize, x: &O, l: usize, r: usize, k: usize) {
        self.eval(l, r, k);
        if b <= l || r <= a {
            return;
        } else if a <= l && r <= b {
            self.lzy[k] = self.lzy[k].op(x);
            self.need_update[k] = true;
            self.eval(l, r, k);
        } else {
            let mid = (l + r) >> 1;
            self._update(a, b, x, l, mid, k << 1);
            self._update(a, b, x, mid, r, (k << 1) + 1);
            self.dat[k] = self.dat[k << 1].op(&self.dat[(k << 1) + 1]);
        }
    }

    pub fn query(&mut self, a: usize, b: usize) -> M {
        let sz = self.size;
        return self._query(a, b, 0, sz, 1);
    }

    pub fn _query(&mut self, a: usize, b: usize, l: usize, r: usize, k: usize) -> M {
        self.eval(l, r, k);
        if b <= l || r <= a {
            return M::id();
        } else if a <= l && r <= b {
            return self.dat[k].clone();
        } else {
            let mid = (l + r) >> 1;
            let vl = self._query(a, b, l, mid, k << 1);
            let vr = self._query(a, b, mid, r, (k << 1) + 1);
            return vl.op(&vr);
        }
    }
}

struct SegTree<M: Monoid> {
    dat: Vec<M>,
    size: usize,
}

impl<M: Monoid + Clone> SegTree<M> {
    pub fn new(n: usize) -> SegTree<M> {
        let mut sz = 1;
        while sz < n {
            sz <<= 1;
        }

        SegTree::<M> {
            dat: vec![M::id(); sz * 2],
            size: sz,
        }
    }

    pub fn update(&mut self, k: usize, val: M) {
        let mut k = k + self.size;
        self.dat[k] = val;
        while k > 0 {
            k = k >> 1;
            self.dat[k] = self.dat[k << 1].op(&self.dat[(k << 1) + 1]);
        }
    }

    pub fn query(&self, a: usize, b: usize) -> M {
        let mut lx = M::id();
        let mut rx = M::id();
        let mut l = a + self.size;
        let mut r = b + self.size;
        while l < r {
            if (l & 1) == 1 {
                lx = lx.op(&self.dat[l]);
                l += 1;
            }

            if (r & 1) == 1 {
                r -= 1;
                rx = self.dat[r].op(&rx);
            }

            l >>= 1;
            r >>= 1;
        }

        return lx.op(&rx);
    }
}

fn decode(v: Vec<char>) -> Vec<(char, usize)> {
    let mut res = Vec::new();
    let mut now = v[0];
    let mut cnt = 1;

    for i in 1..v.len() {
        if v[i] == now {
            cnt += 1;
        } else {
            res.push((now, cnt));
            cnt = 1;
            now = v[i];
        }
    }
    res.push((now, cnt));

    return res;
}

type Num = usize;

#[derive(Clone, Copy)]
pub struct ModInt<T: Copy + Clone>(pub T);

impl Add<ModInt<Num>> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn add(self, rhs: ModInt<Num>) -> ModInt<Num> {
        self + rhs.0
    }
}

impl Add<Num> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn add(self, mut rhs: Num) -> ModInt<Num> {
        if rhs >= MOD {
            rhs %= MOD;
        }
        let mut t = rhs + self.0;
        if t >= MOD {
            t = t - MOD;
        }
        ModInt(t)
    }
}

impl Sub<Num> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn sub(self, rhs: Num) -> ModInt<Num> {
        let rhs = if rhs >= MOD { rhs % MOD } else { rhs };
        let value = if self.0 < rhs { self.0 + MOD } else { self.0 };
        ModInt(value - rhs)
    }
}

impl Sub<ModInt<Num>> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn sub(self, rhs: ModInt<Num>) -> ModInt<Num> {
        self - rhs.0
    }
}

impl AddAssign<Num> for ModInt<Num> {
    fn add_assign(&mut self, other: Num) {
        *self = *self + other;
    }
}

impl AddAssign<ModInt<Num>> for ModInt<Num> {
    fn add_assign(&mut self, other: ModInt<Num>) {
        *self = *self + other;
    }
}

impl SubAssign<Num> for ModInt<Num> {
    fn sub_assign(&mut self, other: Num) {
        *self = *self - other;
    }
}

impl SubAssign<ModInt<Num>> for ModInt<Num> {
    fn sub_assign(&mut self, other: ModInt<Num>) {
        *self = *self - other;
    }
}

impl Div<Num> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn div(self, mut rhs: Num) -> ModInt<Num> {
        if rhs >= MOD {
            rhs %= MOD;
        }
        self * ModInt(rhs).pow(MOD - 2)
    }
}

impl Div<ModInt<Num>> for ModInt<Num> {
    type Output = ModInt<Num>;
    fn div(self, rhs: ModInt<Num>) -> ModInt<Num> {
        self / rhs.0
    }
}

impl DivAssign<Num> for ModInt<Num> {
    fn div_assign(&mut self, rhs: Num) {
        *self = *self / rhs
    }
}

impl DivAssign<ModInt<Num>> for ModInt<Num> {
    fn div_assign(&mut self, rhs: ModInt<Num>) {
        *self = *self / rhs
    }
}

impl Mul<ModInt<Num>> for ModInt<Num> {
    type Output = ModInt<Num>;

    fn mul(self, rhs: ModInt<Num>) -> ModInt<Num> {
        self * rhs.0
    }
}

impl Mul<Num> for ModInt<Num> {
    type Output = ModInt<Num>;

    fn mul(self, mut rhs: Num) -> ModInt<Num> {
        if rhs >= MOD {
            rhs %= MOD;
        }
        let t = (self.0 * rhs) % MOD;
        ModInt(t)
    }
}

impl MulAssign<Num> for ModInt<Num> {
    fn mul_assign(&mut self, rhs: Num) {
        *self = *self * rhs;
    }
}

impl MulAssign<ModInt<Num>> for ModInt<Num> {
    fn mul_assign(&mut self, rhs: ModInt<Num>) {
        *self = *self * rhs;
    }
}

impl ModInt<Num> {
    pub fn pow(self, e: Num) -> ModInt<Num> {
        let mut result = ModInt(1);
        let mut cur = self;
        let mut e = e;
        while e > 0 {
            if e & 1 == 1 {
                result *= cur;
            }
            e >>= 1;
            cur *= cur;
        }
        result
    }
}

const MOD: usize = 1e9 as usize + 7;

struct Eratosthenes {
    n: usize,
    prime: Vec<bool>,
}

impl Eratosthenes {
    pub fn new(n: usize) -> Eratosthenes {
        let mut prime = vec![true; n + 1];
        prime[0] = false;
        prime[1] = false;
        let mut i = 2;
        while i * i <= n {
            if prime[i] {
                let mut p = i * 2;
                while p <= n {
                    prime[p] = false;
                    p += i;
                }
            }
            i += 1;
        }

        Eratosthenes { n: n, prime: prime }
    }
}

fn solve() {
    input! {
        n: i64,
        d: i64
    }
    let mut a = vec![];
    let mut mp = HashMap::new();
    for i in 1..=n {
        for j in 1..=n {
            a.push(i * i + j * j);
            let count = mp.entry(i * i - j * j + d).or_insert(0);
            *count += 1 as i64;
        }
    }

    let mut ans = a.iter().map(|i| mp.get(i).unwrap_or(&0)).sum::<i64>();
    println!("{}", ans);
}

fn main() {
    solve();
}

yukicoder

結果

テストケース

コンパイルメッセージ

ソースコード