24.2-前缀和与树状数组(手撕算法篇)

24 阅读7分钟

24_2-前缀和与树状数组(手撕算法篇)

前言

我们了解了传统的前缀和数组的优缺点以及树状数组解决了传统前缀和数组的那些问题,也了解了差分数组与前缀和数组之间的关系,接下来就通过几道算法题来提升巩固一下这些知识吧。

正餐

1109. 航班预订统计

解题思路

我们先来理解一下题目的意思,其实题目的意思说白了,就是让我们在firsti~lasti这个区间内的数都加上seatsi,然后我们再将每一个数字累加起来。这个其实就是一个区间修改的操作,根据我们基础知识中学到的差分数组的概念,在原数组中的区间操作,就相当于在差分数组中的两次单点操作,即我们需要在差分数组的第firsti位和第lasti+1位分别加上和减去seati。这样,我们就得到了最终统计好的差分数组。题目要统计每个航班的作为总数,那么,根据差分数组与前缀和数组的概念,我们只需要对差分数组求取前缀和即可。

代码演示

/*
 * @lc app=leetcode.cn id=1109 lang=typescript
 *
 * [1109] 航班预订统计
 *
 * https://leetcode-cn.com/problems/corporate-flight-bookings/description/
 *
 * algorithms
 * Medium (57.85%)
 * Likes:    289
 * Dislikes: 0
 * Total Accepted:    61.3K
 * Total Submissions: 105.6K
 * Testcase Example:  '[[1,2,10],[2,3,20],[2,5,25]]\n5'
 *
 * 这里有 n 个航班,它们分别从 1 到 n 进行编号。
 * 
 * 有一份航班预订表 bookings ,表中第 i 条预订记录 bookings[i] = [firsti, lasti, seatsi] 意味着在从
 * firsti 到 lasti (包含 firsti 和 lasti )的 每个航班 上预订了 seatsi 个座位。
 * 
 * 请你返回一个长度为 n 的数组 answer,里面的元素是每个航班预定的座位总数。
 * 
 * 
 * 
 * 示例 1:
 * 
 * 
 * 输入:bookings = [[1,2,10],[2,3,20],[2,5,25]], n = 5
 * 输出:[10,55,45,25,25]
 * 解释:
 * 航班编号        1   2   3   4   5
 * 预订记录 1 :   10  10
 * 预订记录 2 :       20  20
 * 预订记录 3 :       25  25  25  25
 * 总座位数:      10  55  45  25  25
 * 因此,answer = [10,55,45,25,25]
 * 
 * 
 * 示例 2:
 * 
 * 
 * 输入:bookings = [[1,2,10],[2,2,15]], n = 2
 * 输出:[10,25]
 * 解释:
 * 航班编号        1   2
 * 预订记录 1 :   10  10
 * 预订记录 2 :       15
 * 总座位数:      10  25
 * 因此,answer = [10,25]
 * 
 * 
 * 
 * 
 * 提示:
 * 
 * 
 * 1 <= n <= 2 * 10^4
 * 1 <= bookings.length <= 2 * 10^4
 * bookings[i].length == 3
 * 1 <= firsti <= lasti <= n
 * 1 <= seatsi <= 10^4
 * 
 * 
 */

// @lc code=start
class FenwickTree {
    // 用于计算lowbit值
    static lowbit(x) {
        return x & -x;
    }
    // 树状数组
    private c: number[];
    // 数组的下标上线,一般是原数组长度加1
    private size: number;
    constructor(size: number) {
        this.size = size;
        this.c = new Array(size);
        this.c.fill(0);
    }
    /**
     * 往原始数组的第i位增加元素x
     * @param i 
     * @param x 
     */
    public add(i: number, x: number): void {
        while(i <= this.size) {
            // console.log('-->',i, this.c[i], x);
            this.c[i] += x;
            i += FenwickTree.lowbit(i);
        }
    }
    /**
     * 查询原数组前i项和
     * @param i 
     * @returns 
     */
    public query(i: number): number {
        let sum = 0;
        // S[i] = S[i - lowbit(i)] + C[i]
        let idx = i;
        // console.log(this.c);
        while(i) {
            sum += this.c[i];
            i -= FenwickTree.lowbit(i);
            console.log(idx, sum, i);
        }
        return sum;
    }
    /**
     * 根据索引查找原始数组每一项的值
     * @param idx 
     * @returns 
     */
    public at(idx: number): number {
        return this.query(idx) - this.query(idx - 1);
    }
    /**
     * 将元素组第 idx 位的值改成 val
     * @param idx 
     * @param val 
     */
    public update(idx: number, val: number): void {
        console.log(`将第${idx}位的值改成: ${val}`)
        this.add(idx, val - this.at(idx));
    }
    public output(): void {
        let line1 = '';
        let line2 = '';
        let line3 = '';
        let line4 = '';
        let line5 = '';
        this.c.forEach((item, idx) => {
            if(idx === 0) return;
            line1+=String(idx).padStart(6, ' ');
            line2+=String("=").padStart(6, '=');
            line3+=String(item).padStart(6, ' ');
            line4+=String("=").padStart(6, '=');
            line5+=String(this.at(idx)).padStart(6, ' ');
        });
        // 编号
        console.log(line1);
        console.log(line2);
        // 树状数组值,当前所有元素都为1时,树状数组中第i位的值实际上就是lowbit(i)
        console.log(line3);
        console.log(line4);
        // 原数组的值
        console.log(line5+"\n\n");
    }
}
function corpFlightBookings(bookings: number[][], n: number): number[] {
    // 根据我们之前了解的差分数组与前缀和数组的关系,前缀和数组的区间修改等于查分数组的两次单点修改
    // 由于涉及到原数组的修改操作,使用树状数组,我们先初始化一个长度为n+1的树状数组
    const tree = new FenwickTree(n+1);
    // 树状数组中存的值实际上是原数组的查分数组
    bookings.forEach(item => {
        tree.add(item[0], item[2]);
        tree.add(item[1] + 1, -item[2]);
    })
    // 最终的结果实际上是差分数组的前缀和
    const res: number[] = [];
    for(let i=1;i<=n;i++) {
        res.push(tree.query(i));
    }
    return res;
};
// @lc code=end


当然,用上面的树状数组实现比较好理解,也比较好扩展,但看似效率不高,其实我们可以同样使用差分数组和前缀和数组的概念,用更高效的方式解决

function corpFlightBookings(bookings: number[][], n: number): number[] {
    let res: number[] = new Array(n);
    res.fill(0);
    bookings.forEach((item, idx) => {
        res[item[0] - 1] += item[2];
        if(item[1] < n) res[item[1]] -= item[2];
    });
    for(let i=1;i<n;i++) res[i] += res[i-1];
    return res;
};

307. 区域和检索 - 数组可修改

解题思路

这题实际上就是树状数组的裸题,话不多说,直接上代码吧

代码演示

/*
 * @lc app=leetcode.cn id=307 lang=typescript
 *
 * [307] 区域和检索 - 数组可修改
 *
 * https://leetcode-cn.com/problems/range-sum-query-mutable/description/
 *
 * algorithms
 * Medium (53.54%)
 * Likes:    316
 * Dislikes: 0
 * Total Accepted:    25.3K
 * Total Submissions: 47.6K
 * Testcase Example:  '["NumArray","sumRange","update","sumRange"]\n[[[1,3,5]],[0,2],[1,2],[0,2]]'
 *
 * 给你一个数组 nums ,请你完成两类查询,其中一类查询要求更新数组下标对应的值,另一类查询要求返回数组中某个范围内元素的总和。
 * 
 * 实现 NumArray 类:
 * 
 * 
 * 
 * 
 * NumArray(int[] nums) 用整数数组 nums 初始化对象
 * void update(int index, int val) 将 nums[index] 的值更新为 val
 * int sumRange(int left, int right) 返回子数组 nums[left, right] 的总和(即,nums[left] +
 * nums[left + 1], ..., nums[right])
 * 
 * 
 * 
 * 
 * 示例:
 * 
 * 
 * 输入:
 * ["NumArray", "sumRange", "update", "sumRange"]
 * [[[1, 3, 5]], [0, 2], [1, 2], [0, 2]]
 * 输出:
 * [null, 9, null, 8]
 * 
 * 解释:
 * NumArray numArray = new NumArray([1, 3, 5]);
 * numArray.sumRange(0, 2); // 返回 9 ,sum([1,3,5]) = 9
 * numArray.update(1, 2);   // nums = [1,2,5]
 * numArray.sumRange(0, 2); // 返回 8 ,sum([1,2,5]) = 8
 * 
 * 
 * 
 * 
 * 提示:
 * 
 * 
 * 1 
 * -100 
 * 0 
 * -100 
 * 0 
 * 最多调用 3 * 10^4 次 update 和 sumRange 方法
 * 
 * 
 * 
 * 
 */

// @lc code=start
class FenwickTree {
    // 用于计算lowbit值
    static lowbit(x) {
        return x & -x;
    }
    // 树状数组
    private c: number[];
    // 数组的下标上线,一般是原数组长度加1
    private size: number;
    constructor(size: number) {
        this.size = size;
        this.c = new Array(size);
        this.c.fill(0);
    }
    /**
     * 往原始数组的第i位增加元素x
     * @param i 
     * @param x 
     */
    public add(i: number, x: number): void {
        while(i <= this.size) {
            // console.log('-->',i, this.c[i], x);
            this.c[i] += x;
            i += FenwickTree.lowbit(i);
        }
    }
    /**
     * 查询原数组前i项和
     * @param i 
     * @returns 
     */
    public query(i: number): number {
        let sum = 0;
        // S[i] = S[i - lowbit(i)] + C[i]
        let idx = i;
        // console.log(this.c);
        while(i) {
            sum += this.c[i];
            i -= FenwickTree.lowbit(i);
            // console.log(idx, sum, i);
        }
        return sum;
    }
    /**
     * 根据索引查找原始数组每一项的值
     * @param idx 
     * @returns 
     */
    public at(idx: number): number {
        return this.query(idx) - this.query(idx - 1);
    }
    /**
     * 将元素组第 idx 位的值改成 val
     * @param idx 
     * @param val 
     */
    public update(idx: number, val: number): void {
        console.log(`将第${idx}位的值改成: ${val}`)
        this.add(idx, val - this.at(idx));
    }
    public output(): void {
        let line1 = '';
        let line2 = '';
        let line3 = '';
        let line4 = '';
        let line5 = '';
        this.c.forEach((item, idx) => {
            if(idx === 0) return;
            line1+=String(idx).padStart(6, ' ');
            line2+=String("=").padStart(6, '=');
            line3+=String(item).padStart(6, ' ');
            line4+=String("=").padStart(6, '=');
            line5+=String(this.at(idx)).padStart(6, ' ');
        });
        // 编号
        console.log(line1);
        console.log(line2);
        // 树状数组值,当前所有元素都为1时,树状数组中第i位的值实际上就是lowbit(i)
        console.log(line3);
        console.log(line4);
        // 原数组的值
        console.log(line5+"\n\n");
    }
}
class NumArray {
    tree: FenwickTree;
    constructor(nums: number[]) {
        this.tree = new FenwickTree(nums.length+1);
        nums.forEach((item, idx) => this.tree.add(idx+1, item));
    }

    update(index: number, val: number): void {
        this.tree.add(index + 1, val - this.tree.at(index+1));
    }

    sumRange(left: number, right: number): number {
        return this.tree.query(right + 1) - this.tree.query(left);
    }
}

/**
 * Your NumArray object will be instantiated and called as such:
 * var obj = new NumArray(nums)
 * obj.update(index,val)
 * var param_2 = obj.sumRange(left,right)
 */
// @lc code=end


面试题 10.10. 数字流的秩

解题思路

首先,我们来翻译一下这道题的意思,说白了,就是我们会不断的读入一些数字,并在读入的过程中,需要查询某一个数字之前出现小于或等于这个数字的数量。我们先来拆解一下上面的问题,首先,我们要读入数字,然后需要计数,然后需要查询不大于某个数字的数字总数。

  1. 读入数字,我们既然之后需要计算数字的数量,那么,我们就可以在读入数字时,使用类似之前学习的计数排序的方式,在一个数组中,数组的索引代表读入的数字,而数组中存储的值代表该数字出现的次数。

    # 读入数字
    1
    # 计数
    [0,1]
    # 读入数字
    5
    # 计数
    [0,1,0,0,0,1]
    # 读入数字
    1
    [0,2,0,0,0,1]
    
  2. 大家可以发现,我们再读入数据的过程中,就已经借助计数数组的方式对每一个数字进行计数了

  3. 接下来就要解决,如果查询不大于某个数字的数字个数呢?比如说我们想要查询不小于3的数字个数,大家观察一下上面的数组就会发现,直接求取索引为3的位置之前的前缀和即可,即不小于3的数字个数为2。由于我们的数字是不断读入的,计数数组也在不断的变化,因此,我们选择使用树状数组去计算前缀和。

代码演示

class FenwickTree {
    // 用于计算lowbit值
    static lowbit(x) {
        return x & -x;
    }
    // 树状数组
    private c: number[];
    // 数组的下标上线,一般是原数组长度加1
    private size: number;
    constructor(size: number) {
        this.size = size;
        this.c = new Array(size);
        this.c.fill(0);
    }
    /**
     * 往原始数组的第i位增加元素x
     * @param i 
     * @param x 
     */
    public add(i: number, x: number): void {
        while(i <= this.size) {
            // console.log('-->',i, this.c[i], x);
            this.c[i] += x;
            i += FenwickTree.lowbit(i);
        }
    }
    /**
     * 查询原数组前i项和
     * @param i 
     * @returns 
     */
    public query(i: number): number {
        let sum = 0;
        // S[i] = S[i - lowbit(i)] + C[i]
        let idx = i;
        // console.log(this.c);
        while(i) {
            sum += this.c[i];
            i -= FenwickTree.lowbit(i);
            // console.log(idx, sum, i);
        }
        return sum;
    }
    /**
     * 根据索引查找原始数组每一项的值
     * @param idx 
     * @returns 
     */
    public at(idx: number): number {
        return this.query(idx) - this.query(idx - 1);
    }
    /**
     * 将元素组第 idx 位的值改成 val
     * @param idx 
     * @param val 
     */
    public update(idx: number, val: number): void {
        console.log(`将第${idx}位的值改成: ${val}`)
        this.add(idx, val - this.at(idx));
    }
    public output(): void {
        let line1 = '';
        let line2 = '';
        let line3 = '';
        let line4 = '';
        let line5 = '';
        this.c.forEach((item, idx) => {
            if(idx === 0) return;
            line1+=String(idx).padStart(6, ' ');
            line2+=String("=").padStart(6, '=');
            line3+=String(item).padStart(6, ' ');
            line4+=String("=").padStart(6, '=');
            line5+=String(this.at(idx)).padStart(6, ' ');
        });
        // 编号
        console.log(line1);
        console.log(line2);
        // 树状数组值,当前所有元素都为1时,树状数组中第i位的值实际上就是lowbit(i)
        console.log(line3);
        console.log(line4);
        // 原数组的值
        console.log(line5+"\n\n");
    }
}
class StreamRank {
    private tree: FenwickTree;
    constructor() {
        // 依据题目条件,x最大为50000,由于树状数组第0位是没有实际用处的,我们加上一个偏移量,如2
        this.tree = new FenwickTree(50002);
    }

    track(x: number): void {
        // 插入时x加上偏移量1
        this.tree.add(x+1, 1);
    }

    getRankOfNumber(x: number): number {
        // 查询时x加上偏移量1
        return this.tree.query(x + 1)
    }
}

/**
 * Your StreamRank object will be instantiated and called as such:
 * var obj = new StreamRank()
 * obj.track(x)
 * var param_2 = obj.getRankOfNumber(x)
 */