Ссылка на задачу с LeetCode — 2338. Count the Number of Ideal Arrays.
📋 Описание задачи
Дан массив длины n
, состоящий из целых чисел от 1
до maxValue
. Массив называется «идеальным», если:
- Каждый элемент принадлежит диапазону
[1, maxValue]
- Каждый следующий элемент делится на предыдущий
Нужно посчитать количество различных идеальных массивов длины n
.
Ответ вернуть по модулю 10⁹ + 7
.
💬 Пример
- Вход:
n = 2, maxValue = 5
- Выход:
10
- Пояснение: Возможные идеальные массивы:
- Начинаются с
1
(5
массивов): [1,1], [1,2], [1,3], [1,4], [1,5]
- Начинаются с
2
(2
массива): [2,2], [2,4]
- Начинаются с
3
(1
массив): [3,3]
- Начинаются с
4
(1
массив): [4,4]
- Начинаются с
5
(1
массив): [5,5]
- Всего:
5 + 2 + 1 + 1 + 1 = 10
идеальных массивов.
💡 Идея
- Любую цепочку
i₀ → i₁ → ... → iₙ
можно поделить на i₀
, сведя задачу к подсчёту цепочек, начинающихся с единицы:
1 → j₁ → j₂ → ... → jₙ = j
- Рассмотрим разложение
j
на простые множители:
j = p₁^a₁ * p₂^a₂ * ... * pₖ^aₖ
- Для каждого простого
pᵢ
, мы должны распределить его aᵢ
вкладов по n - 1
стрелкам между числами в цепочке.
- Это задача о разбиении на суммы, и методом шаров и перегородок показывается, что это число равно:
C(n−1+aᵢ,aᵢ)
- Общее число способов построить цепочку — это произведение таких значений по всем простым множителям
j
.
🛠 Подробности реализации
- Предварительно вычисляем биномиальные коэффициенты
C(n - 1 + k, k)
через модульную арифметику с обратными по модулю.
- Генерируем решето Эратосфена до
maxValue
для быстрого разложения на простые множители.
- Для каждого
val
от 1
до maxValue
:
- Разлагаем
val
на простые множители.
- Для каждой степени
aᵢ
считаем C(n - 1 + aᵢ, aᵢ)
и перемножаем.
- Умножаем результат на количество кратных
val
значений (maxValue / val
).
- Складываем все такие значения и возвращаем по модулю
10⁹ + 7
.
⏱ Асимптотика
- Время:
O(V log V)
, где V = maxValue
- Решето:
O(V log log V)
- Цикл по всем значением с разложением на множители каждого за
O(log val)
- Память:
O(V)
- Решето простых:
O(V)
- Биномиальный кэш:
O(log V)
🧾 Исходный код
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