главная новое лучшее написать
3

Ссылка на задачу с LeetCode — 2338. Count the Number of Ideal Arrays.

📋 Описание задачи

Дан массив длины n, состоящий из целых чисел от 1 до maxValue. Массив называется «идеальным», если:

Нужно посчитать количество различных идеальных массивов длины n.
Ответ вернуть по модулю 10⁹ + 7.

💬 Пример

💡 Идея

🛠 Подробности реализации

  1. Предварительно вычисляем биномиальные коэффициенты C(n - 1 + k, k) через модульную арифметику с обратными по модулю.
  2. Генерируем решето Эратосфена до maxValue для быстрого разложения на простые множители.
  3. Для каждого val от 1 до maxValue:
    • Разлагаем val на простые множители.
    • Для каждой степени aᵢ считаем C(n - 1 + aᵢ, aᵢ) и перемножаем.
    • Умножаем результат на количество кратных val значений (maxValue / val).
  4. Складываем все такие значения и возвращаем по модулю 10⁹ + 7.

⏱ Асимптотика

🧾 Исходный код

struct N;

impl Solution {
    pub fn ideal_arrays(n: i32, max_value: i32) -> i32 {
        const MOD: i32 = 1_000_000_007;
        let n = n as usize;
        let max_value = max_value as usize;

        // Precompute dp[k] = C(n - 1 + k, k) using modular inverse
        let dp = Self::precompute_sum_splits(n - 1, N::MAX_BITS, MOD);

        let mut result = 0i64;
        for val in 1..=max_value {
            // Number of ideal arrays with base `val`
            let ways = Self::count_ideal_arrays_with_base(val, MOD as i64, &dp);
            let multiples = (max_value / val) as i64;
            result = (result + multiples * ways) % MOD as i64;
        }

        result as i32
    }

    // Factorizes `base` and multiplies corresponding dp[k] for each prime power
    fn count_ideal_arrays_with_base(mut base: usize, modulo: i64, dp: &[i32]) -> i64 {
        let mut result = 1i64;
        while base > 1 {
            let prime = N::SIEVE[base];
            let mut count = 0;
            while base % prime == 0 {
                base /= prime;
                count += 1;
            }
            result = result * dp[count] as i64 % modulo;
        }
        result
    }

    // Computes dp[k] = C(n - 1+ k, k) for k = 1..=max_k using modular inverse
    // C(n - 1 + k, k) = number of ways to insert k prime factors into n slots (with repetition)
    fn precompute_sum_splits(n: usize, max_k: usize, modulo: i32) -> Vec<i32> {
        let modulo = modulo as i64;
        let mut dp = vec![0; max_k + 1];
        dp[0] = 1; // not used, but aligns indexing
        for k in 1..=max_k {
            let num = dp[k - 1] as i64 * (n - 1 + k) as i64 % modulo;
            let denom_inv = N::modinv(k as i64, modulo);
            dp[k] = (num * denom_inv % modulo) as i32;
        }
        dp
    }
}

impl N {
    const MAX_VALUE: usize = 10_000;
    const MAX_BITS: usize = Self::msb(Self::MAX_VALUE);
    const SIEVE: [usize; Self::MAX_VALUE + 1] = Self::generate_sieve();

    // Computes modular inverse of n modulo p
    fn modinv(n: i64, p: i64) -> i64 {
        Self::powmod(n, p - 2, p)
    }

    // Fast exponentiation (n^k mod p)
    fn powmod(mut base: i64, mut exp: i64, modulo: i64) -> i64 {
        let mut result = 1;
        while exp > 0 {
            if exp % 2 == 1 {
                result = result * base % modulo;
            }
            base = base * base % modulo;
            exp >>= 1;
        }
        result
    }

    // Computes smallest prime factor for all numbers ≤ MAX_VALUE
    const fn generate_sieve() -> [usize; Self::MAX_VALUE + 1] {
        let mut spf = [0; Self::MAX_VALUE + 1];
        let mut p = 2;
        while p <= Self::MAX_VALUE {
            if spf[p] == 0 {
                let mut multiple = p;
                while multiple <= Self::MAX_VALUE {
                    spf[multiple] = p;
                    multiple += p;
                }
            }
            p += 1;
        }
        spf
    }

    // Computes number of significant bits for n - binary logarithm
    const fn msb(n: usize) -> usize {
        if n <= 1 { 1 } else { 1 + Self::msb(n / 2) }
    }
}

Tags: #rust #algorithms #combinatorics