Сегодня чуток похулиганим и предоставим переоптимизированное решение для задачи 2554. Maximum Number of Integers to Choose From a Range I. В реальном интервью такое могут потребовать лишь на уровне идеи, но нам интересно запрограммировать самим :)
Описание задачи 📋
Необходимо выбрать максимально возможное количество чисел из диапазона [1,n], при этом соблюдая ограничения:
- Числа из списка banned выбирать нельзя.
- Каждое число можно использовать не более одного раза.
- Сумма выбранных чисел не должна превышать maxSum.
Результатом должно быть количество чисел, которые можно выбрать, удовлетворяя этим условиям.
Идея решения 💡
Предвычисляем суммы запрещённых чисел и используем двоичный поиск, чтобы найти максимальное k, для которого сумма допустимых чисел в диапазоне [1,k] не превышает maxSum. Это позволяет эффективно учитывать ограничения и избежать лишних вычислений.
Обзор решения 🧠
-
Сортировка и удаление дубликатов:
- Сортируем массив
banned и удаляем повторяющиеся элементы.
-
Предвычисление кумулятивных сумм:
- Создаем массив
ban_sums, где на позиции i содержится сумма первых i запрещённых чисел. Это позволяет быстро вычислять сумму запрещённых чисел до любого предела.
-
Двоичный поиск:
- Выполняем двоичный поиск по диапазону
[1,n], чтобы найти максимальное k, при котором сумма чисел в диапазоне [1,k] (с учётом исключений) не превышает maxSum.
- Для каждого шага поиска вычисляем сумму через поиск в массиве
ban_sums за O(log B).
-
Окончательный подсчёт:
- От найденного
k вычитаем количество запрещённых чисел в диапазоне [1,k], чтобы получить итоговое количество допустимых чисел.
Анализ сложности 📊
-
Временная сложность:
- Сортировка массива banned:
O(B log B).
- Двоичный поиск по n:
O(log N), где каждая проверка требует O(log B).
- Итоговая сложность:
O(log B⋅(B + log B)).
-
Пространственная сложность:
O(B) для хранения массива ban_sums.
Замечение по реализации
К сожалению, стандартная библиотека Rust предоставляет бинарный поиск только в векторах и слайсах, а значит реализацию для range'ей предстоит написать самим.
Исходный код решения
use std::cmp::Ordering;
impl Solution {
pub fn max_count(banned: Vec<i32>, n: i32, max_sum: i32) -> i32 {
// Step 1: Sort and remove duplicates from the banned list
let mut banned = banned;
banned.sort_unstable();
banned.dedup();
// Step 2: Precompute cumulative sums for banned numbers
let mut ban_sums = vec![0i64; banned.len() + 1];
for (i, &v) in banned.iter().enumerate() {
ban_sums[i + 1] = ban_sums[i] + v as i64;
}
// Step 3: Use binary_search_by to find the largest k such that the sum is <= max_sum
let max_sum = max_sum as i64;
let k = Self::binary_search_in_range(0..=n, |k| {
let (_, sum) = Self::calculate_sum(k, &banned, &ban_sums);
sum.cmp(&max_sum)
}).unwrap_or_else(|bound| bound - 1);
// Step 4: Calculate the result based on the largest valid k
let (n_banned, _) = Self::calculate_sum(k, &banned, &ban_sums);
k - n_banned as i32
}
/// Helper function to calculate number of banned elements <= k and the sum
fn calculate_sum(k: i32, banned: &Vec<i32>, ban_sums: &Vec<i64>) -> (usize, i64) {
let n_banned = banned.binary_search(&(k + 1)).unwrap_or_else(|x| x);
let total_sum = (k as i64 * (k as i64 + 1)) / 2;
let adjusted_sum = total_sum - ban_sums[n_banned];
(n_banned, adjusted_sum)
}
/// Helper function: Binary search on a range using a custom condition.
/// Returns `Ok(value)` if found, else `Err(bound)` where `bound` is the first greater element.
fn binary_search_in_range<F>(range: std::ops::RangeInclusive<i32>, condition: F) -> Result<i32, i32>
where
F: Fn(i32) -> Ordering,
{
let mut left = *range.start();
let mut right = *range.end();
while left <= right {
let mid = left + (right - left) / 2;
match condition(mid) {
Ordering::Less => left = mid + 1,
Ordering::Equal => return Ok(mid),
Ordering::Greater => right = mid - 1,
}
}
Err(left)
}
}