https://atcoder.jp/contests/abc272/tasks/abc272_d
素因数分解さえできれば、移動できるステップを高速に求められます。そのために、GaussIntという構造体を作りました。
デバッグのためにfmtというメソッドを作って複素数を出力できるようにしましたが、その中で
write!(f, "{}", self.re)?;
などとしていますが、?はwrite!でエラーが出たときはそのエラーをそのまま返せるようにする記号です。
この規模だと速くはなりませんね。
#![allow(non_snake_case)]
use std::ops::{Add, Mul};
use std::collections::VecDeque;
use std::fmt;
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()
}
fn read_input() -> (i32, i32) {
let v = read_vec();
(v[0], v[1])
}
fn div_pow(n: i32, d: i32) -> (i32, i32) {
if n % d == 0 {
let (e, m) = div_pow(n / d, d);
(e + 1, m)
}
else {
(0, n)
}
}
type Factors = Vec<(i32, i32)>;
fn factorize(n: i32, p0: i32) -> Factors {
if n == 1 {
return vec![]
}
for p in p0.. {
if p * p > n {
break
}
if n % p == 0 {
let (e, m) = div_pow(n, p);
let factors = factorize(m, p + 1);
let mut new_factors = vec![(p, e)];
new_factors.extend(&factors);
return new_factors
}
}
vec![(n, 1)]
}
fn pow_mod(n: i32, e: i32, d: i32) -> i32 {
if e == 1 {
n
}
else if e % 2 == 0 {
let m = pow_mod(n, e / 2, d);
((m as i64) * (m as i64) % (d as i64)) as i32
}
else {
let m = pow_mod(n, e - 1, d);
((m as i64) * (n as i64) % (d as i64)) as i32
}
}
#[derive(Clone, Copy)]
#[derive(Debug)]
#[derive(Eq, Hash, PartialEq)]
struct Point {
x: i32,
y: i32,
}
impl Add for Point {
type Output = Self;
fn add(self, other: Self) -> Self {
Self { x: self.x + other.x, y: self.y + other.y }
}
}
fn rotate(dir: &Point) -> Point {
Point { x: -dir.y, y: dir.x }
}
fn reflect(dir: &Point) -> Point {
Point { x: dir.y, y: dir.x }
}
fn expand_dirs(dir: Point) -> Vec<Point> {
let mut dirs = Vec::<Point>::new();
dirs.push(dir);
for i in 0..3 {
dirs.push(rotate(&dirs[i]))
}
if dir.x != dir.y && dir.x != 0 && dir.y != 0 {
for i in 0..4 {
dirs.push(reflect(&dirs[i]))
}
}
dirs
}
#[derive(Clone, Copy)]
struct GaussInt {
re: i32,
im: i32
}
impl Mul for GaussInt {
type Output = Self;
fn mul(self, other: Self) -> Self {
let re = self.re * other.re - self.im * other.im;
let im = self.re * other.im + self.im * other.re;
Self { re, im }
}
}
impl fmt::Display for GaussInt {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
if self.re != 0 {
write!(f, "{}", self.re)?;
if self.im > 0 {
write!(f, "+")?;
}
}
else if self.im == 0 {
write!(f, "0")?
}
if self.im == 1 {
write!(f, "i")?
}
else if self.im == -1 {
write!(f, "-i")?
}
else if self.im == 0 {
return Ok(())
}
else {
write!(f, "{}i", self.im)?
};
return Ok(())
}
}
impl GaussInt {
fn conjugate(&self) -> GaussInt {
GaussInt { re: self.re, im: -self.im }
}
fn rotate90(&self) -> GaussInt {
GaussInt { re: -self.im, im: self.re }
}
fn is_normal(&self) -> bool {
self.re > 0 && self.re > self.im.abs()
}
fn normalize(&self) -> GaussInt {
let mut c = *self;
for _ in 0..4 {
if c.is_normal() {
return c
}
c = c.rotate90()
}
c
}
fn find_i(p: i32) -> i32 {
let e = p - 1;
for a in 2..p {
let k = pow_mod(a, e / 2, p);
if k == p - 1 {
return pow_mod(a, e / 4, p)
}
}
1
}
fn gcd(z1: &GaussInt, z2: &GaussInt) -> GaussInt {
if z2.im == 0 {
return *z1
}
else {
let z3 = z2.normalize();
let q = z1.re / z3.re;
let r = z1.re % z3.re;
let z4 = GaussInt { re: r, im: z1.im - q * z3.im };
GaussInt::gcd(&z3, &z4)
}
}
fn pow(&self, e: i32) -> GaussInt {
if e == 0 {
GaussInt { re: 1, im: 0 }
}
else if e == 1 {
*self
}
else if e % 2 == 0 {
let z = self.pow(e / 2);
z * z
}
else {
*self * self.pow(e - 1)
}
}
fn divide_prime(p: i32) -> GaussInt {
let k = GaussInt::find_i(p);
let z1 = GaussInt { re: p, im: 0 };
let z2 = GaussInt { re: k, im: 1 };
let z = GaussInt::gcd(&z1, &z2);
GaussInt { re: z.re, im: z.im.abs() }
}
}
fn powers(z: GaussInt, e: i32, num_e: i32) -> Vec<GaussInt> {
let zc = z.conjugate();
(0..num_e).map(|i| z.pow(e-i) * zc.pow(i)).collect()
}
fn expand_zs(zs0: Vec<GaussInt>, (p, e): (i32, i32),
head: bool) -> (Vec<GaussInt>, bool) {
let (ne, new_head) = if p == 2 { (1, head) }
else if head { (e/2 + 1, false) }
else { (e+1, false) };
let zs = if p == 2 { vec![(GaussInt { re: 1, im: 1 }).pow(e)] }
else if p % 4 == 1 { powers(GaussInt::divide_prime(p), e, ne) }
else { vec![GaussInt { re: p.pow((e/2) as u32), im: 0 }] };
(zs0.iter().map(|&z0| zs.iter().map(|&z| z0*z).collect::<Vec<GaussInt>>()).
flatten().collect(), new_head)
}
fn find_normal_dirs(factors: &Factors, first: bool) -> Vec<GaussInt> {
let zs0 = vec![GaussInt { re: 1, im: 0 }];
let p = factors.iter().fold((zs0, first),
|(zs, head), &fs| expand_zs(zs, fs, head));
p.0
}
fn find_movable_dirs(M: i32) -> Vec<Point> {
let factors = factorize(M, 2);
if factors.iter().any(|&(p, e)| p % 4 == 3 && e % 2 == 1) {
return vec![]
}
let zs = find_normal_dirs(&factors, true);
zs.into_iter().map(|z| expand_dirs(Point { x: z.re, y: z.im })).
flatten().collect::<Vec<Point>>()
}
fn make_table(N: i32, M: i32) -> Vec<Vec<i32>> {
let mut table: Vec<Vec<i32>>
= (0..N).map(|_| (0..N).map(|_| -1).collect()).collect();
table[0][0] = 0;
let dirs = find_movable_dirs(M);
let mut queue = VecDeque::<Point>::new();
queue.push_back(Point { x: 0, y: 0 });
while queue.len() > 0 {
let pt0 = queue.pop_front().unwrap();
let x0 = pt0.x as usize;
let y0 = pt0.y as usize;
for dir in dirs.iter() {
let pt = pt0 + *dir;
let x = pt.x as usize;
let y = pt.y as usize;
if 0 <= pt.x && pt.x < N && 0 <= pt.y && pt.y < N &&
table[x][y] == -1 {
table[x][y] = table[x0][y0] + 1;
queue.push_back(pt)
}
}
}
table
}
fn print_table(table: Vec<Vec<i32>>) {
for v in table.into_iter() {
println!("{}", v.iter().map(|n| n.to_string()).
collect::<Vec<String>>().join(" "))
}
}
fn f(N: i32, M: i32) {
let table = make_table(N, M);
print_table(table)
}
fn main() {
let v = read_input();
f(v.0, v.1)
}
