▲

3

▼

💡 Эффективный подсчет с использованием битовых операций

leetcoder, 13-12-2024 , 158

Очередная наша задача - 2593. Find Score of an Array After Marking All Elements. Ну и так как наша цель не просто решать, а показывать имплементацию разных техник, то воспользуемся-ка мы идеей битового множества (его в стандартную библиотеку Rust еще "не подвезли", так что и операции имплементируем сами).

Задача 🧩

Дан массив положительных целых чисел nums. Ваша цель — получить максимальное значение score, выполняя следующие действия:

1. Выберите наименьший немаркированный элемент массива (при равенстве — элемент с меньшим индексом).
2. Добавьте значение этого элемента к score.
3. Пометьте выбранный элемент и его двух соседей (если они существуют).
4. Повторяйте, пока все элементы не будут помечены.

Необходимо вернуть итоговое значение score.

Идея 💡

В данной задаче нужно эффективно следовать указанным правилам. Чтобы сделать это:

• (а) Мы отсортируем индексы массива nums по значениям, чтобы обрабатывать элементы в правильном порядке, начиная с наименьших.
• (б) Используем битовое множество (Vec<u128>) для компактного представления состояния "помечен" для элементов массива.

Подход 🚀

1. Сортировка индексов: Сортируем индексы массива nums по значениям элементов, чтобы начинать с наименьших чисел.
2. Использование Vec<u128>:
   • Каждый бит в массиве Vec<u128> представляет состояние одного элемента массива.
   • Для упрощения проверки границ смещаем индексы на 1.
3. Обновление состояния:
   • Для текущего элемента и его соседей выставляем соответствующие биты в Vec<u128> с помощью битовых операций.
   • Оптимизируем выставление сразу нескольких битов одной операцией в наиболее частом случае, когда индексы находятся внутри границ одной u128 ячейки.
   • Аккуратно обрабатываем пограничные случаи, учитывая переход между блоками u128.
4. Подсчет score: Добавляем значение текущего элемента к итоговому результату, если он еще не помечен.

Сложность ⏱️

• Временная сложность:
  • Сортировка занимает O(nlog⁡n).
  • Обработка каждого элемента: O(n).
  • Итого: O(nlog⁡n).
• Пространственная сложность:
  • O(n/128) для Vec<u128>.
  • O(n) для хранения отсортированных индексов.
  • Итого: O(n).

Исходный код решения

impl Solution {
    pub fn find_score(nums: Vec<i32>) -> i64 {
        let n = nums.len();
        let mut indices: Vec<usize> = (0..n).collect();

        // Sort indices based on the values in nums, using sort_by_key for simplicity
        indices.sort_by_key(|&idx| nums[idx]);

        // Use Vec<u128> to mark elements and their adjacent ones,
        // add reserved places to skip edge cases checks
        let mut marked_bits = vec![0u128; (n + 129) / 128];
        let mut score: i64 = 0;

        for &idx in &indices {
            // Adjust the index by adding 1 to simplify boundary conditions
            let adjusted_idx = idx + 1;
            let bit_offset = adjusted_idx % 128;
            let block_idx = adjusted_idx / 128;

            // Check if the current index is already marked
            if ((marked_bits[block_idx] >> bit_offset) & 1) != 0 {
                continue;
            }

            // Add the value at nums[idx] to the score
            score += nums[idx] as i64;

            if 0 < bit_offset && bit_offset < 127 {
                // Mark the adjacent elements in the common case
                marked_bits[block_idx] |= 0b101 << (bit_offset - 1);
            } else {
                // Handle edge cases for elements at the boundaries
                if bit_offset == 0 {
                    marked_bits[block_idx - 1] |= 1 << 127;
                    marked_bits[block_idx] |= 1 << 1;
                } else if bit_offset == 127 {
                    marked_bits[block_idx] |= 1 << 126;
                    marked_bits[block_idx + 1] |= 1;
                }
            }
        }

        score
    }
}