1508. 子数组和排序后的区间和

题目描述

leetcode-cn.com/problems/ra…

分析

如果用暴力法求出所有区间和，排序后再做统计，那么一定会超时，所以得想个投机取巧的方法

仔细去审视上述过程，可以发现，如果在计算区间和过程中，完成排序而不是所有的都算完再排，不就行了吗

所以很显然，可以借助堆来进行这个操作

算法

堆

过程

核心是在统计区间和的时候，自动完成排序～

堆中元素是什么？

首先，堆中保存的元素是一种特殊的数据结构，{ i, j, val }，其中：

i -> 区间左边界 j -> 区间右边界 val -> 区间和

初始化

数组中的所有位置都是我的区间的 i, 因此所有的区间，一开始都是 { i, i, nums[i] }

通过遍历不断进行一下操作：

• extract
• 更新统计结果
• insert

其中，需要注意以下几点：

统计的时机

遍历到 left - 1, 因为编号从 1 开始

insert 时机

item.j + 1 合法，即 item.j + 1 < n

代码

/**
 * @param {number[]} nums
 * @param {number} n
 * @param {number} left
 * @param {number} right
 * @return {number}
 */
var rangeSum = function(nums, n, left, right) {
    let ans = 0
    const h = new Heap((a, b) => {
        if (!b) return false
        return a.val > b.val
    })
    for (let i = 0; i < n; i++) {
        h.insert(getData(i, i, nums[i]))
    }
    
    for (let i = 0; i < right; i++) {
        const item = h.extract()
        if (i >= left - 1) {
            ans = (ans + item.val) % 1000000007
        }
        if (item.j + 1 < n) h.insert(getData(item.i, item.j + 1, nums[item.j + 1] + item.val))
    }
    
    return ans
};

function getData(i, j, val) {
    return { i, j, val }
}

class Heap {
    constructor(cmp) {
        this.cmp = cmp
        this.heap = []
    }
    
    getLeftIndex(index) {
        return index * 2 + 1
    }
    
    getRightIndex(index) {
        return index * 2 + 2
    }
    
    getParentIndex(index) {
        return Math.floor((index - 1) / 2)
    }
    
    size() {
        return this.heap.length
    }
    
    swap(parent, index) {
        const arr = this.heap
        ;[arr[parent], arr[index]] = [arr[index], arr[parent]]
    }
    
    isEmpty() {
        return this.size() === 0
    }
    
    insert(value) {
        const index = this.heap.length
        this.heap.push(value)
        this.siftUp(index)
    }
    
    siftUp(index) {
        let parent = this.getParentIndex(index)
        
        while (index > 0 && this.cmp(this.heap[parent], this.heap[index])) {
            this.swap(parent, index)
            index = parent
            parent = this.getParentIndex(index)
        }
    }
    
    extract() {
        if (this.isEmpty()) return
        if (this.size() === 1) return this.heap.pop()
        
        const removedItem = this.heap[0]
        this.heap[0] = this.heap.pop()
        this.siftDown(0)
        
        return removedItem
    }
    
    siftDown(index) {
        let element = index
        const left = this.getLeftIndex(index)
        const right = this.getRightIndex(index)
        
        if (index < this.size() && this.cmp(this.heap[element], this.heap[left])) {
            element = left
        }
        
        if (index < this.size() && this.cmp(this.heap[element], this.heap[right])) {
            element = right
        }
        
        if (element !== index) {
            this.swap(element, index)
            this.siftDown(element)
        }
    }
    
    top() {
        return this.heap[0]
    }
}