AtCoder Grand Contest 038 C - LCMs

問題

atcoder.jp

解法メモ

 \rm lcm(A_i, A_j) を求めるのではなく \frac{A_i A_j}{\rm gcd(A_i,A_j)} を求めていく方針を考える。

サンプル2の以下の例を解くことを考える。

8
1 2 3 4 6 8 12 12
 \sum _{ i \lt j } A_i A_j = (\sum_{i} A_i )^2 - \sum_{i} A_i^2

なので、i j のペアについて考えるのではなく単体で考えてあとで調整できるということを頭に留めておく。

最大公約数ごとにグルーピングしたりしたいが、とりあえずそのあたりは大変そうなので、まずは約数ごとにまとめて眺めてみる。

1 => {1, 2, 3, 4, 6, 8, 12, 12}
2 => {2, 4, 6, 8, 12, 12}
3 => {3, 6, 12, 12}
4 => {4, 8, 12, 12}
6 => {6, 12, 12}
8 => {8}
12 => {12, 12}

これについて、最大公約数はとりあえず置いておいて、約数で割って計算するとどうなるか見てみる。

 \frac{12 \times 12}{12} + \frac{12 \times 12 + 12 \times 6 + 12 \times 6}{6} + ...

という感じになる。この部分で、 \frac{12 \times 12}{12} + \frac{12 \times 6 + 12 \times 6}{6} + ...は最終的な回答に含まれるが、 \frac{12 \times 12}{6} は含まれないので取り除かなければならない。

なので、最終的な答えについて考えてみると

 \frac{12 \times 12}{12} + \frac{12 \times 12 + 12 \times 6 + 12 \times 6}{6} + ...  - \frac{12 \times 12}{6} - \frac{12 \times 12}{4} - \frac{12 \times 12}{3} - \frac{12 \times 12}{2} - \frac{12 \times 12}{1}

という感じになる。

コード

use self::mod_int::ModInt;
const MOD: usize = 998244353;
const N: usize = 1000001;

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());

    let mut divisors = vec![vec![]; N];
    for i in 1..N {
        let mut cur = i;
        while cur < N {
            divisors[cur].push(i);
            cur += i;
        }
    }

    let mut inv = vec![ModInt(1); N];
    for i in 1..N {
        inv[i] /= i;
    }

    let n: usize = sc.read();
    let mut count = vec![0; N];
    for _ in 0..n {
        let a: usize = sc.read();
        count[a] += 1;
    }

    let mut sum = vec![ModInt(0); N];
    let mut sq_sum = vec![ModInt(0); N];
    for i in 0..N {
        if count[i] == 0 {
            continue;
        }

        let count = count[i];
        for &divisor in divisors[i].iter() {
            sum[divisor] += ModInt(i) * count;
            sq_sum[divisor] += ModInt(i) * i * count;
        }
    }

    let mut sum_by_d = vec![ModInt(0); N];
    for d in (1..N).rev() {
        sum_by_d[d] += (sum[d] * sum[d] - sq_sum[d]) * inv[2] * inv[d];
        for &divisor in divisors[d].iter() {
            if d == divisor {
                continue;
            }
            let sum_by_d_d = sum_by_d[d];
            sum_by_d[divisor] -= sum_by_d_d * d * inv[divisor];
        }
    }

    let mut ans = ModInt(0);
    for d in 1..N {
        ans += sum_by_d[d];
    }

    println!("{}", ans.0);
}

pub mod mod_int {
    use super::MOD;
    use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

    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
        }
    }
}
pub 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 buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    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()
    }
}