AGC 003 D - Anticube

問題

https://beta.atcoder.jp/contests/agc003/tasks/agc003_d

解説

各 s_i を素因数分解する（方法は後で考える）。

まず、 s_i に 3 つ同じ素因数が含まれるとき、それらを取り除いてしまっても構わない。例えば のとき、 を取り除いて と考えて良い。このように各 s_i を素数の 3 乗で割れるだけ割っておく。すると、各 s_i は と表せる。ここで前者を p_i 後者を q_i として とする。この s_i との積が立法数になる s_j は と表せるはずである。

なので、問題は と の数をそれぞれ求め、多い方を採用するようにする。また、 と表せるものは 1 つだけ採用できる。

最初に戻って s_i を素因数分解する方法を考える。 となるような全ての素数 p について調べたあとの s_i は素数の 3 乗を含まないので、以下の 3 つのいずれかである。

• 素数
• 2 つの相異なる素数の積
• 素数の2乗

素数の2乗のとき、 と表せる。そうでない時は と表せる。このレベルまで分かれば十分なので、完全に素因数分解する必要はなく、実行時間制限に間に合う。

コード

use std::cmp;
use std::collections::{BTreeMap, BTreeSet};

fn get_primes(n: usize) -> Vec<u64> {
    let mut is_prime = vec![true; n + 1];
    is_prime[0] = false;
    is_prime[1] = false;
    for p in 2..(n + 1) {
        if !is_prime[p] {
            continue;
        }
        let mut cur = 2 * p;
        while cur <= n {
            is_prime[cur] = false;
            cur += p;
        }
    }

    is_prime
        .iter()
        .enumerate()
        .filter(|&(_, &is_prime)| is_prime)
        .map(|(p, _)| p as u64)
        .collect()
}

fn main() {
    let primes = get_primes(100000);
    let prime2 = primes
        .iter()
        .map(|&p| (p * p, p))
        .collect::<BTreeMap<u64, u64>>();
    let mut sc = Scanner::new();
    let n = sc.usize_read();

    let mut count = BTreeMap::new();

    for _ in 0..n {
        let mut s: u64 = sc.read();
        let mut s1 = 1;
        let mut s2 = 1;
        for &prime in &primes {
            let p2 = prime * prime;
            let p3 = p2 * prime;
            if p3 > s {
                break;
            }
            while s % p3 == 0 {
                s /= p3;
            }
            if s % p2 == 0 {
                s /= p2;
                s2 *= prime;
            } else if s % prime == 0 {
                s /= prime;
                s1 *= prime;
            }
        }

        match prime2.get(&s) {
            Some(&p) => s2 *= p,
            None => s1 *= s,
        }

        let &cur = count.get(&(s1, s2)).unwrap_or(&0);
        count.insert((s1, s2), cur + 1);
    }

    let mut ans = 0;
    let mut set = BTreeSet::new();
    set.insert((1, 1));
    for &(s1, s2) in count.keys() {
        if set.contains(&(s1, s2)) {
            continue;
        }
        let &count1 = count.get(&(s1, s2)).unwrap_or(&0);
        let &count2 = count.get(&(s2, s1)).unwrap_or(&0);
        ans += cmp::max(count1, count2);
        set.insert((s1, s2));
        set.insert((s2, s1));
    }

    if count.contains_key(&(1, 1)) {
        ans += 1;
    }
    println!("{}", ans);
}

struct Scanner {
    ptr: usize,
    length: usize,
    buf: Vec<u8>,
    small_cache: Vec<u8>,
}

#[allow(dead_code)]
impl Scanner {
    fn new() -> Scanner {
        Scanner {
            ptr: 0,
            length: 0,
            buf: vec![0; 1024],
            small_cache: vec![0; 1024],
        }
    }

    fn load(&mut self) {
        use std::io::Read;
        let mut s = std::io::stdin();
        self.length = s.read(&mut self.buf).unwrap();
    }

    fn byte(&mut self) -> u8 {
        if self.ptr >= self.length {
            self.ptr = 0;
            self.load();
            if self.length == 0 {
                self.buf[0] = b'\n';
                self.length = 1;
            }
        }

        self.ptr += 1;
        return self.buf[self.ptr - 1];
    }

    fn is_space(b: u8) -> bool {
        b == b'\n' || b == b'\r' || b == b'\t' || b == b' '
    }

    fn read_vec<T>(&mut self, n: usize) -> Vec<T>
    where
        T: std::str::FromStr,
        T::Err: std::fmt::Debug,
    {
        (0..n).map(|_| self.read()).collect()
    }

    fn usize_read(&mut self) -> usize {
        self.read()
    }

    fn read<T>(&mut self) -> T
    where
        T: std::str::FromStr,
        T::Err: std::fmt::Debug,
    {
        let mut b = self.byte();
        while Scanner::is_space(b) {
            b = self.byte();
        }

        for pos in 0..self.small_cache.len() {
            self.small_cache[pos] = b;
            b = self.byte();
            if Scanner::is_space(b) {
                return String::from_utf8_lossy(&self.small_cache[0..(pos + 1)])
                    .parse()
                    .unwrap();
            }
        }

        let mut v = self.small_cache.clone();
        while !Scanner::is_space(b) {
            v.push(b);
            b = self.byte();
        }
        return String::from_utf8_lossy(&v).parse().unwrap();
    }
}