Это решение задачи 4 Leetcode 98 двухнедельных конкурсов. О других задачах с предыдущих конкурсов можно прочитать здесь. Кажется, это самая сложная задача в этом конкурсе.

Проблема

Вам даны два массива 0-indexed nums1 и nums2 и двумерный массив запросов queries. Существует три типа запросов:

Для запроса типа 1 queries[i] = [1, l, r]. Переверните значения с 0 на 1 и с 1 на 0 в nums1 с индекса l на индекс r. И l, и r имеют индекс 0.

Для запроса типа 2 queries[i] = [2, p, 0]. Для каждого индекса 0 <= i < n установите nums2[i] = nums2[i] + nums1[i] * p.

Для запроса типа 3 queries[i] = [3, 0, 0]. Найдите сумму элементов в nums2.

Возвращает массив, содержащий все ответы на запросы третьего типа.

Решение

Если вам нужно ответить на несколько запросов о сумме/минимум/максимуме диапазона массива, а также обновить значения массива, вы должны подумать, подходит ли задача для структуры данных Дерево сегментов. Существует несколько оптимизаций дерева сегментов. В нашем случае нам нужно обновлять диапазоны, а не отдельные элементы. Вот специальная оптимизация дерева сегментов, предназначенная для обновления диапазона. Также похоже, что нам не нужно немедленно применять каждое обновление. Давайте рассмотрим запрос первого типа. Если мы применим его четное количество раз к одному и тому же диапазону элементов, результат будет таким же, как если бы мы не применяли запросы к этим элементам. Так что, похоже, в нашем случае лучше откладывать обновления в дереве сегментов как можно дольше, пока они нам действительно не понадобятся. К счастью, существует специальная оптимизация дерева сегментов, называемая ленивым распространением. Также давайте посмотрим на операцию обновления. Это не обычное обновление по некоторым номерам. Это флип 0–1. К счастью, существует известная реализация дерева сегментов, рассчитанная на (-1) и 1 флип. Подробнее можно узнать здесь.

Подход

Хорошо, давайте представим, что у нас есть дерево сегментов с ленивым распространением, предназначенное для флипов со следующим интерфейсом:

class SegmentTree {
  // Builds segment tree from array
  build(arr: Array<number>);
  // Flips elements in the array from start index to end index
  update(start: number, end: nuymber);
  // Return sum of elements from start index to end index
  querySum(start: number, end: number);
}

Сделаем следующее. Обновляем массив nums1. Мы заменяем каждый 0 в нем на -1, чтобы можно было использовать листание дерева сегментов. Вычислим начальную сумму массива nums2 и построим дерево отрезков из массива nums1. Для каждого запроса:

  • Если тип запроса = 1. Мы вызываем обновление для дерева сегментов, чтобы перевернуть элементы из index query[1] в query[2]
  • Если тип запроса = 2, мы обновляем предварительно вычисленную сумму. Для этого нам нужно знать число 1 в массиве nums1. Предполагая, что мы изменили каждый 0 на -1, число 1 можно рассчитать следующим образом: number_of_ones=(array_sum + array_length)/2. Изменение суммы можно рассчитать как new_sum=sum + number_of_ones * запрос[1].
  • Если тип запроса = 3, просто поместите предварительно вычисленную сумму в массив результатов.
var handleQuery = function(nums1, nums2, queries) {
    // We change all 0 in nums1 to (-1) to allow usage of flipping segment tree
    for(let i = 0; i < n; i++){
        if(nums1[i] === 0){
            nums1[i] = -1;
        }
    }
    const tree = new SegmentTree();
    // Build segment tree out of array
    tree.build(nums1);
    let sum = 0;
    // Calculate initial sum of nums2 array
    for(let num of nums2){
        sum+= num
    }
    let res = []
    for(let q of queries){
        // If query type is 1
        if(q[0] === 1){
            // Flip elements in array
            tree.update(q[1], q[2])
        }
        if(q[0] === 2){
            // We update array sum to new one. Sum change could be calculated
            // as following = number_of_ones * query_multiplier
            // Where number_of_ones = (array_sum + array_length)/2
            sum += ((n + query_tree(1, 0, n-1, 0, n-1))/2 * q[1])
        }
        if(q[0] === 3) {
            // Just move precalculated array sum to result array
            res.push(sum)
        }
    }
    
    return res
};

Сложность

Временная сложность:O(q*log(n) + n*log(n)), где q — количество запросов, а n — количество элементов.

Пространственная сложность:O(4*n)

P.S.

Вот мое решение на Leetcode.