結果

問題 No.3333 Consecutive Power Sum (Large)
コンテスト
ユーザー akakimidori
提出日時 2025-11-02 22:32:44
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 3,447 ms / 10,000 ms
コード長 8,086 bytes
コンパイル時間 11,350 ms
コンパイル使用メモリ 399,760 KB
実行使用メモリ 60,308 KB
最終ジャッジ日時 2025-11-02 22:33:56
合計ジャッジ時間 65,325 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 63
権限があれば一括ダウンロードができます
コンパイルメッセージ
warning: unused import: `std::io::Write`
  --> src/main.rs:22:5
   |
22 | use std::io::Write;
   |     ^^^^^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` on by default

warning: type alias `Map` is never used
  --> src/main.rs:24:6
   |
24 | type Map<K, V> = BTreeMap<K, V>;
   |      ^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: type alias `Set` is never used
  --> src/main.rs:25:6
   |
25 | type Set<T> = BTreeSet<T>;
   |      ^^^

warning: type alias `Deque` is never used
  --> src/main.rs:26:6
   |
26 | type Deque<T> = VecDeque<T>;
   |      ^^^^^

ソースコード

diff # 

// E=1 の列挙はどうやる？
// 素因数分解できるなら前と同じやつやればいい
// まあ1個くらいなら間に合うだろう
//
// 尺取りは可能？
// E=2は無理
// E=3 はギリギリ間に合うかというくらい
// 間に合いそうなので考えない
//
// E=2はどうする？
// R(R+1)(2R+1)/6 - (L-1)L(2L-1)/6 = N
// L^2+(L+1)^2+..+R^2
// 因数分解できる？
//
// a = R - L + 1
// b = L + R
// として
// a(a^2 + 3b^2 - 1) / 12
// これも素因数分解してaを試してbを固定してみる？

use std::collections::*;
use std::io::Write;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn main() {
    /*
    for l in 1u128..=100 {
        for r in l..=100 {
            let s = (l..=r).map(|i| i.pow(2)).sum::<u128>();
            let n = r - l + 1;
            let k = l + r;
            let v = n * (n * n + 3 * k * k - 1);
            assert_eq!(v, 12 * s, "{l} {r}");
        }
    }
    */
    input!(n: u128);
    let mut ans = vec![];
    // E=1
    let m = 2 * n;
    for d in divisor(m) {
        // R-L+1 = d
        // L + R = m/d
        let x = d - 1;
        let y = m / d;
        if (x + y) % 2 != 0 || x >= y {
            continue;
        }
        let r = (x + y) / 2;
        let l = (y - x) / 2;
        if l <= r && 1 <= l {
            ans.push((1, l, r));
        }
    }
    // E = 2
    let m = 12 * n;
    for d in divisor(m) {
        let q = m / d;
        if d.saturating_mul(d) - 1 >= q {
            continue;
        }
        let b = q - d * d + 1;
        if b % 3 != 0 {
            continue;
        }
        let b = b / 3;
        let sq = kth_root(b, 2);
        if sq * sq != b {
            continue;
        }
        let a = d - 1;
        let b = sq;
        if a & 1 != b & 1 || a >= b {
            continue;
        }
        // a = R - L + 1
        // b = L + R
        let l = (b - a) / 2;
        let r = (a + b) / 2;
        ans.push((2, l, r));
    }
    for e in (3..).take_while(|i| 1u128 << i <= n) {
        let mut s = 0;
        let mut l = 1u128;
        for r in (1u128..).take_while(|i| i.pow(e) <= n) {
            s += r.pow(e);
            while s > n {
                s -= l.pow(e);
                l += 1;
            }
            if s == n {
                ans.push((e, l, r));
            }
        }
    }
    ans.sort();
    println!("{}", ans.len());
    for (e, l, r) in ans {
        println!("{e} {l} {r}");
    }
}

pub fn divisor(n: u128) -> Vec<u128> {
    let f = factorize(n);
    let mut d = vec![1];
    for (p, c) in f {
        let l = d.len();
        for _ in 0..c {
            let k = d.len();
            for i in 0..l {
                let v = d[k - l + i] * p;
                d.push(v);
            }
        }
    }
    d
}

// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

#[macro_export]
macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

#[macro_export]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------

// n = prod_i p_i ^ a_i
// (p_i, a_i) の列を返す
pub fn factorize(n: u128) -> Vec<(u128, u32)> {
    assert!(n > 0);
    let mut list = vec![];
    let mut stack = vec![];
    stack.push(n);
    while let Some(n) = stack.pop() {
        if n <= 1 {
            continue;
        }
        if is_prime_miller(n) {
            list.push((n, 1));
            continue;
        }
        if n % 2 == 0 {
            list.push((2, 1));
            stack.push(n / 2);
            continue;
        }
        'pollard: for i in 1u128.. {
            let f = |x: u128| (mul(x, x, n) + i) % n;
            let mut x = i;
            let mut y = f(x);
            loop {
                let g = binary_gcd(y - x + n, n);
                if g == 0 || g == n {
                    break;
                }
                if g != 1 {
                    stack.push(g);
                    stack.push(n / g);
                    break 'pollard;
                }
                x = f(x);
                y = f(f(y));
            }
        }
    }
    list.sort();
    list.dedup_by(|a, b| (a.0 == b.0).then(|| b.1 += a.1).is_some());
    list
}

// ---------- begin miller-rabin ----------
pub fn is_prime_miller(n: u128) -> bool {
    if n <= 1 {
        return false;
    } else if n <= 3 {
        return true;
    } else if n % 2 == 0 {
        return false;
    }
    let pow = |r: u128, mut m: u128| -> u128 {
        let mut t = 1u128;
        let mut s = r % n;
        let n = n as u128;
        while m > 0 {
            if m & 1 == 1 {
                t = mul(t, s, n);
            }
            s = mul(s, s, n);
            m >>= 1;
        }
        t
    };
    let mut d = n - 1;
    let mut s = 0;
    while d % 2 == 0 {
        d /= 2;
        s += 1;
    }
    const B: [u128; 13] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41];
    for &b in B.iter() {
        let mut a = pow(b, d);
        if a <= 1 {
            continue;
        }
        let mut i = 0;
        while i < s && a != n - 1 {
            i += 1;
            a = mul(a, a, n);
        }
        if i >= s {
            return false;
        }
    }
    true
}
// ---------- end miller-rabin ----------

pub fn mul(a: u128, b: u128, m: u128) -> u128 {
    let mut x = [0; 3];
    let mut y = [0; 3];
    for i in 0..3 {
        x[i] = a >> (30 * i) & ((1 << 30) - 1);
        y[i] = b >> (30 * i) & ((1 << 30) - 1);
    }
    let mut ans = x[2] * y[2];
    ans = ((ans << 30) + (x[1] * y[2] + x[2] * y[1])) % m;
    for i in (0..3).rev() {
        ans <<= 30;
        for k in 0..=i {
            ans += x[k] * y[i - k];
        }
        ans %= m;
    }
    /*
    let mut ans = 0;
    for i in (0..90).rev() {
        ans = ans + ans;
        if ans >= m {
            ans -= m;
        }
        if b >> i & 1 == 1 {
            ans += a;
        }
    }
    */
    ans
}

// ---------- begin binary_gcd ----------
pub fn binary_gcd(a: u128, b: u128) -> u128 {
    if a == 0 || b == 0 {
        return a + b;
    }
    let x = a.trailing_zeros();
    let y = b.trailing_zeros();
    let mut a = a >> x;
    let mut b = b >> y;
    while a != b {
        let x = (a ^ b).trailing_zeros();
        if a < b {
            std::mem::swap(&mut a, &mut b);
        }
        a = (a - b) >> x;
    }
    a << x.min(y)
}
// ---------- end binary_gcd ----------
// floor(a^(1/k))
pub fn kth_root(a: u128, k: u64) -> u128 {
    assert!(k > 0);
    if a == 0 {
        return 0;
    }
    if k >= 128 {
        return 1;
    }
    if k == 1 {
        return a;
    }
    let mut v = ((a as f64).ln() / k as f64).exp().round() as u128;
    while v.checked_pow(k as u32).map_or(true, |p| p > a) {
        v -= 1;
    }
    while (v + 1).checked_pow(k as u32).map_or(false, |p| p <= a) {
        v += 1;
    }
    v
}
0