AtCoder Beginner Contest 446 E

https://atcoder.jp/contests/abc446/tasks/abc446_e

3項の漸化式なので、前後のペアをノードとします。M²通りあります。そのノードを辿っていき、0が含まれずにループになったら条件を満たすノードです。

// Many LCMs
#![allow(non_snake_case)]

use std::collections::HashSet;


//////////////////// 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() -> (usize, usize, usize) {
    let v: Vec<usize> = read_vec();
    let (M, A, B) = (v[0], v[1], v[2]);
    (M, A, B)
}

type Point = (usize, usize);

fn F_each(mut pt: Point, M: usize, A: usize, B: usize,
                            memo: &Vec<Vec<i32>>) -> (Vec<Point>, i32) {
    let mut path: Vec<(usize, usize)> = vec![pt];
    let mut visited: HashSet<(usize, usize)> = HashSet::new();
    loop {
        let (x, y) = pt;
        if x == 0 || y == 0 {
            return (path, 0)
        }
        else if memo[x][y] != -1 {
            return (path, memo[x][y])
        }
        else if visited.contains(&pt) {
            return (path, 1)
        }
        else {
            visited.insert(pt);
            pt = ((x*A + y*B) % M, x);
            path.push(pt)
        }
    }
}

fn F(M: usize, A: usize, B: usize) -> i32 {
    // -1: 未確定 0: 含む 1: 含まない
    let mut memo: Vec<Vec<i32>> = vec![vec![-1; M]; M];
    for x in 0..M {
        for y in 0..M {
            let (path, c) = F_each((x, y), M, A, B, &memo);
            for (x1, y1) in path {
                memo[x1][y1] = c
            }
        }
    }
    memo.into_iter().map(|v| v.into_iter().sum::<i32>()).sum::<i32>()
}

fn main() {
    let (M, A, B) = read_input();
    println!("{}", F(M, A, B))
}