Neliöt

warning: unused variable: `j`
  --> input/code.rs:21:9
   |
21 |     for j in (0..).map(|j| 1 << j).take_while(|j| *j < even) {
   |         ^ help: consider prefixing with an underscore: `_j`
   |
   = note: `#[warn(unused_variables)]` on by default

warning: unused variable: `i`
  --> input/code.rs:82:21
   |
82 |                 for i in (0..).take_while(|i| *i * *i <= n) {
   |                     ^ help: consider prefixing with an underscore: `_i`

use rand::prelude::*;
use std::collections::BTreeMap;
use std::io::BufRead;

fn powmod(b: u64, e: u64, m: u64) -> u64 {
    if e == 0 {
        1
    } else if e % 2 == 0 {
        let p = powmod(b, e / 2, m) as u128;
        ((p * p) % m as u128) as u64
    } else {
        ((b as u128 * powmod(b, e - 1, m) as u128) % m as u128) as u64
    }
}

fn is_w(a: u64, even: u64, odd: u64, p: u64) -> bool {
    let mut u = powmod(a, odd, p);
    if u == 1 {
        return false;
    }
    for j in (0..).map(|j| 1 << j).take_while(|j| *j < even) {
        if u == p - 1 {
            return false;
        }
        u = ((u as u128 * u as u128) % p as u128) as u64;
    }
    true
}

fn is_prime(p: u64) -> bool {
    if p == 2 {
        true
    } else if p <= 1 || p % 2 == 0 {
        false
    } else {
        let mut odd = p - 1;
        let mut even = 1;
        while odd % 2 == 0 {
            even *= 2;
            odd /= 2;
        }
        for b in &[2, 325, 9375, 28178, 450775, 9780504, 1795265022] {
            let a = b % p;
            if a == 0 {
                return true;
            }
            if is_w(a, even, odd, p) {
                return false;
            }
        }
        true
    }
}

pub fn gcd(a: u64, b: u64) -> u64 {
    if b == 0 {
        a
    } else {
        gcd(b, a % b)
    }
}

fn factor(n: u64) -> BTreeMap<u64, usize> {
    fn step(x: &mut u64, n: u64, c: u64) {
        *x = ((*x as u128 * *x as u128 + c as u128) % n as u128) as u64;
    }

    fn f(n: u64, r: &mut BTreeMap<u64, usize>) {
        let mut n = n;
        while n % 2 == 0 {
            n /= 2;
            *r.entry(2).or_insert(0) += 1;
        }
        if n == 1 {
        } else if is_prime(n) {
            *r.entry(n).or_insert(0) += 1;
        } else {
            loop {
                let mut x = thread_rng().next_u64() % n;
                let mut y = x;
                let c = thread_rng().next_u64() % n;
                for i in (0..).take_while(|i| *i * *i <= n) {
                    step(&mut x, n, c);
                    step(&mut x, n, c);
                    step(&mut y, n, c);
                    let g = gcd(x.max(y) - x.min(y), n);
                    if g == n {
                        break;
                    } else if g > 1 {
                        f(n / g, r);
                        f(g, r);
                        return;
                    }
                }
            }
        }
    }

    let mut res = BTreeMap::new();
    if n > 1 {
        f(n, &mut res);
    }
    res
}

fn solve(c: usize) -> bool {
    for (p, k) in factor(c as u64) {
        if p % 4 == 3 && k % 2 == 1 {
            return false;
        }
    }
    true
}

fn main() {
    let stdin = std::io::stdin();
    let stdin = stdin.lock();
    let mut lines = stdin.lines();

    let t: usize = lines.next().unwrap().unwrap().parse().unwrap();
    for _ in 0..t {
        let c: usize = lines.next().unwrap().unwrap().parse().unwrap();
        if solve(c) {
            println!("YES");
        } else {
            println!("NO");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_prime() {
        assert!(is_prime(2));
        assert!(is_prime(3));
        assert!(is_prime(5));
        assert!(is_prime(7));
        assert!(is_prime(11));
        assert!(is_prime(13));
        assert!(is_prime(1_000_000_007));
        assert!(is_prime(1_000_000_009));

        assert!(!is_prime(14));
        assert!(!is_prime(15));
        assert!(!is_prime(1_000_000_007 * 1_000_000_009));
    }

    #[test]
    fn test_factor() {
        assert_eq!(factor(16), [(2, 4)].into_iter().copied().collect());
        assert_eq!(factor(15), [(3, 1), (5, 1)].into_iter().copied().collect());
        assert_eq!(
            factor(1_000_000_007 * 1_000_000_009),
            [(1_000_000_007, 1), (1_000_000_009, 1)]
                .into_iter()
                .copied()
                .collect()
        );
    }
}

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: ACCEPTED

Test 3

Group: 2

Verdict: ACCEPTED