#595187 (Rust) No.1322 Totient Bound

• 問題: No.1322 Totient Bound
• ユーザー: akakimidori
• 提出日時: 2020-12-19 01:58:52
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 739 ms / 5,000 ms
• コード長: 8,663 bytes
• コンパイル時間: 13,065 ms
• コンパイル使用メモリ: 378,960 KB
• 実行使用メモリ: 61,472 KB
• 最終ジャッジ日時: 2024-09-21 09:49:35
• 合計ジャッジ時間: 25,604 ms
• ジャッジサーバーID （参考情報）: judge3 / judge2

ファイルパターン	結果
sample	AC * 3
other	AC * 36

ソースコード

// 昔の提出から
// https://atcoder.jp/contests/xmascon19/submissions/11229505

struct FenwickTree {
    val: Vec<u64>,
}
 
impl FenwickTree {
    fn new(n: usize) -> Self {
        FenwickTree {
            val: vec![0; n + 1],
        }
    }
    fn add(&mut self, mut x: usize, v: u64) {
        while let Some(p) = self.val.get_mut(x) {
            *p = p.wrapping_add(v);
            x += x & (!x + 1);
        }
    }
    fn sum(&mut self, mut x: usize) -> u64 {
        assert!(x < self.val.len());
        let mut ans = 0u64;
        while x > 0 {
            ans = ans.wrapping_add(self.val[x]);
            x -= x & (!x + 1);
        }
        ans
    }
}
 
fn sieve(n: usize) -> Vec<u64> {
    let mut is_prime = vec![true; n + 1];
    for i in 2.. {
        if i * i > n {
            break;
        }
        if is_prime[i] {
            let mut j = i * i;
            while j <= n {
                is_prime[j] = false;
                j += i;
            }
        }
    }
    let len = is_prime.iter().skip(2).filter(|p| **p).count();
    let mut prime = Vec::with_capacity(len);
    for (i, is_prime) in is_prime.into_iter().enumerate().skip(2) {
        if is_prime {
            prime.push(i as u64);
        }
    }
    prime
}
 
struct PrimeCountSolver {
    query: Vec<u64>,
    memo: Vec<u64>,
}
 
impl PrimeCountSolver {
    fn new() -> Self {
        PrimeCountSolver {
            query: vec![],
            memo: vec![],
        }
    }
    fn add(&mut self, n: u64) {
        self.query.push(n);
    }
    fn build(&mut self) {
        self.query.sort();
        self.query.dedup();
        let n = self.query.last().map_or(1, |n| *n);
        let mut m = 1;
        while (m + 1) * (m + 1) <= n {
            m += 1;
        }
        let p = sieve(m as usize);
        let bound = (n as f64).cbrt().powi(2) as u64 + 2;
        let mut stack = vec![];
        for &v in self.query.iter() {
            let k = match p.binary_search_by(|&p| (p * p).cmp(&v)) {
                Ok(k) => k + 1,
                Err(k) => k,
            };
            stack.push((v, k));
        }
        let mut query = vec![];
        while let Some((n, k)) = stack.pop() {
            if k == 0 {
                continue;
            }
            let q = p[k - 1];
            if q * q > n {
                let x = match p[..k].binary_search_by(|&p| (p * p).cmp(&n)) {
                    Ok(k) => k + 1,
                    Err(k) => k,
                };
                stack.push((n, x));
            } else if n <= bound {
                query.push((k, n));
            } else {
                stack.push((n, k - 1));
                stack.push((n / q, k - 1));
            }
        }
        query.sort();
        query.dedup();
        let m = bound as usize;
        let mut bit = FenwickTree::new(m);
        for i in 1..(m + 1) {
            bit.add(i, 1);
        }
        let mut is_prime = vec![true; m + 1];
        let mut memo = vec![0; query.len()];
        let mut pos = 0;
        for (i, p) in p.iter().enumerate() {
            let p = *p as usize;
            let mut j = p;
            while j <= m {
                if is_prime[j] {
                    bit.add(j, !0);
                }
                is_prime[j] = false;
                j += p;
            }
            while let Some(&(k, n)) = query.get(pos) {
                if k > i + 1 {
                    break;
                } else {
                    memo[pos] = bit.sum(n as usize);
                    pos += 1;
                }
            }
            if pos >= query.len() {
                break;
            }
        }
        self.memo.clear();
        self.memo.resize(self.query.len(), 0);
        let mut stack = vec![];
        for (i, &n) in self.query.iter().enumerate() {
            let k = match p.binary_search_by(|&p| (p * p).cmp(&n)) {
                Ok(k) => k + 1,
                Err(k) => k,
            };
            self.memo[i] += k as u64 - 1;
            stack.push((n, k, 1, i));
        }
        while let Some((n, k, sign, i)) = stack.pop() {
            if k == 0 {
                self.memo[i] += sign * n;
                continue;
            }
            let q = p[k - 1];
            if q * q > n {
                let x = match p[..k].binary_search_by(|p| (p * p).cmp(&n)) {
                    Ok(k) => k + 1,
                    Err(k) => k,
                };
                self.memo[i] += (1 ^ !sign) * (k - x) as u64;
                stack.push((n, x, sign, i));
            } else if n <= bound {
                self.memo[i] += sign * memo[query.binary_search(&(k, n)).unwrap()];
            } else {
                stack.push((n, k - 1, sign, i));
                stack.push((n / q, k - 1, 1 ^ !sign, i));
            }
        }
    }
    fn get(&self, n: u64) -> u64 {
        self.memo[self.query.binary_search(&n).unwrap()]
    }
}
// 昔の提出終わり

// ---------- begin enumerate prime ----------
fn enumerate_prime(n: usize) -> Vec<usize> {
    assert!(n <= 10usize.pow(8));
    if n <= 1 {
        return vec![];
    }
    let batch = (n as f64).sqrt() as usize + 1;
    let mut is_prime = vec![true; batch + 1];
    for i in (2..).take_while(|p| p * p <= batch) {
        if is_prime[i] {
            let mut j = i * i;
            while let Some(p) = is_prime.get_mut(j) {
                *p = false;
                j += i;
            }
        }
    }
    let mut prime = vec![];
    let mut small_prime = vec![];
    for (i, p) in is_prime.iter().enumerate().skip(2) {
        if *p && i <= n {
            prime.push(i);
            small_prime.push(i);
        }
    }
    let mut l = batch;
    while l <= n {
        let r = std::cmp::min(l + batch, n + 1);
        is_prime.clear();
        is_prime.resize(r - l, true);
        for &p in small_prime.iter() {
            let mut j = (l + p - 1) / p * p - l;
            while let Some(is_prime) = is_prime.get_mut(j) {
                *is_prime = false;
                j += p;
            }
        }
        for (i, p) in is_prime.iter().enumerate().filter(|p| *p.1) {
            if *p {
                prime.push(i + l);
            }
        }
        l += batch;
    }
    prime
}
// ---------- end enumerate prime ----------

// ?
// 知ってること
// N*2*3/2*5/4*... でテキトーな上限Mが計算できて、
// phi(x) <= N, x <= M なものを数えられればいい
// x = ** p^a
// phi(x) = ** (p-1)*p^(a-1)
// phi(x) <= x
// 素数カウント、phiの和
//
// 最後にかける素数についてのみなら素数カウントが使えないこともない
// それまでの列挙の計算量が爆発してるから厳しいのでは
// よく考えれば考える状態自体はsqrt(N) ではある
// でもsqrt(N) くらい遷移する必要があるからやっぱ厳しいのでは
// もっと早い段階で枝刈りができるんじゃないか？
// できそう
// (p-1)^2 > K なやつはさっさと素数カウントの方に突っ込んでいくと良さげ

#[allow(dead_code)]
fn run(n: usize) {
    let m = (1..).find(|m| (m - 1) * (m - 1) > n).unwrap() - 1;
    let prime = enumerate_prime(m);
    let mut dp = std::collections::BTreeMap::new();
    let mut count = std::collections::BTreeMap::new();
    let mut ans = 0;
    dp.insert(n, 1);
    for (i, &p) in prime.iter().enumerate() {
        let mut memo = vec![];
        for (&n, &c) in dp.iter().rev().take_while(|e| *e.0 >= p - 1) {
            let mut d = p - 1;
            while d <= n {
                memo.push((n / d, c));
                d *= p;
            }
        }
        for (n, v) in memo {
            *dp.entry(n).or_insert(0) += v;
        }
        let mut del = vec![];
        for (&n, &c) in dp.iter() {
            if (p - 1).pow(2) <= n {
                break;
            }
            del.push((n, c));
        }
        for (n, c) in del {
            dp.remove(&n);
            *count.entry(n + 1).or_insert(0) += c;
            let x = prime.binary_search_by_key(&(n + 1, 1), |p| (*p, 0)).unwrap_err();
            ans -= x.min(i + 1) * c;
        }
    }
    for (n, c) in dp {
        *count.entry(n + 1).or_insert(0) += c;
        ans -= prime.len() * c;
    }
    let mut solver = PrimeCountSolver::new();
    for (&n, _) in count.iter() {
        solver.add(n as u64);
    }
    solver.build();
    for (&n, &c) in count.iter() {
        ans += (solver.get(n as u64) as usize + 1) * c;
    }
    println!("{}", ans);
}

fn main() {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).unwrap();
    let n: usize = s.trim().parse().unwrap();
    run(n);
}