Solving Every CSES Problems in Rust - #1 Number Spiral

Problem Constraints

The constraints for x, y are 10^9. Given these constraints, the problem cannot be solved by actually building the spiral - but a static function returns the solution in O(1) / O(log n).

Pattern Recognition

The problem is a simple pattern matching problem. The first observation to be made is that for a given (x, y), the value would be greater than squared(max(x, y) - 1).

Then only 4 cases emerge:

• x > y and square(x) is even → Add y to square(x)
• x > y and square(x) is odd → Add 2*x - y to square(x)
• y > x and square(x) is even → Add 2*y - x to square(x)
• y > x and square(x) is odd → Add x to square(x)

Rust Implementation

use std::io;

fn get_res(x: i64, y: i64, b: bool) -> i64 {
    let mut res = (x - 1).pow(2);
    if (res % 2 == 0) == b {
        res += 2 * x - y;
    } else {
        res += y;
    }
    res
}

fn main() {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    let t: i32 = input.trim().parse().unwrap();

    for _ in 0..t {
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();
        let mut iter = input.split_whitespace();

        let x: i64 = iter.next().unwrap().parse().unwrap();
        let y: i64 = iter.next().unwrap().parse().unwrap();

        let res;
        if x > y {
            res = get_res(x, y, false);
        } else {
            res = get_res(y, x, true);
        }

        println!("{}", res);
    }
}