[路飞]计算右侧小于左侧当前元素的个数-树状数组

95 阅读4分钟

记录 1 道算法题

计算右侧小于左侧当前元素的个数

要求:提供一个整数数组,返回数组中的每一个元素的在他右边且比他小的数的个数。

比如:

输入 [5,2,6,1] 返回 [2,1,1,0]

最简单的解法就是遍历数组的每一个元素,然后双重循环,在里面对后面的元素进行一一对比,将计数存到结果数组的对应位置。

采用算法进行优化的话有两种方法,一种是根据桶和树状数组,另一个是借助归并排序,归并在另一篇文章进行介绍。

桶和树状数组

桶的使用就是根据数据的最大值 n 创造一个从 0 到 n 的数组(如果最小不是 0 的话,为了和下标一一对应还要做偏移计算),这个数组就叫做桶,在遍历过程中对遍历到的数在桶中 + 1,当遍历完之后,桶内记录的就是每个数出现的次数。并且这个桶是升序的。我们可以从后往前遍历数组,这样桶内记录的数就是当前遍历元素的右边的数各自的数量。接下来只要计算比当前元素小的数有多少个就好了。

给定的数据不一定占满整个桶,比如 [1,2,5] ,创建的桶是 [0,0,0,0,0],总共大小是 5 个位置,其中 3 和 4 是浪费空间的,他们的数量永远是 0,这时候有一个离散化的概念。就是将没用上的位置去掉,减少桶的长度,但是需要构建一个下标的映射关系,让数和下标对应起来。比如上面的例子离散化之后的桶是 [0, 0, 0],其中有一个映射关系 { 1: 0, 2: 1, 5: 3 }

如何动态计算每次桶中的数则用上了前缀和,动态维护前缀和的更好的方法是树状数组。但仅从实现上看,我们知道从后往前遍历,每遍历一个数,比这个元素大的数的前缀和 + 1,有多少个增加多少取决于比当前元素大的数有多少个。这实际上是在上一次的桶的计数总和上递增的,所以这是会构成一个前缀和数组。

如果不用树状数组实现的话,流程大概是:

[5,2,6,1]
桶: [0,0,0,0]
映射关系:{
            5: 3,
            2: 2,
            6: 4,
            1: 1
         }
前缀和数组:[0,0,0,0,0] // 多一个是为了方便计数前一个的和, 前缀和数组记录的是小于等于下标 i 元素的数的个数
从后往前遍历,插入结果是在前面推入数组 unshift
而 前缀和数组[i - 1] 就是比当前元素小的个数

当前元素 1
1 往后的前缀和都 + 1
前缀和数组:[0,1,1,1,1]
结果:[0]
当前元素 6
前缀和数组:[0,1,1,1,2]
结果:[1,0]
当前元素 2
前缀和数组:[0,1,2,2,3]
结果:[1,1,0]
当前元素 5
前缀和数组:[0,1,2,3,4]
结果:[2,1,1,0]

上面的思路使用代码实现则是:

    // nums 是传入的数组, map 是映射,
    const result = []
    const b = new Array(nums.length).fill(0) // 发挥了桶和前缀和的作用
    for(let i = nums.length - 1; i >= 0; i--) {
        const n = nums[i]
        for(let j = map[n]; j < b.length; j++) {
            b[j]++
        }
        
        result.unshift(b[map[n] - 1])
    }

当然上面这种普通做法耗时非常长,树状数组通过二分法减少了一些循环递增的过程,下面我们看一下使用树状数组怎么解题。

树状数组采用了二进制进行管理,通过二分法进行操作,计算使用的函数也是著名的 lowbit。想知道更详细的可以在网上搜 lowbit,这里仅贴我看过的两篇文章。

  1. 浅谈lowbit运算 - Seaway-Fu - 博客园 (cnblogs.com),
  2. 树状数组 & lowbit() - liubilan - 博客园 (cnblogs.com)

分为查询和更新两个操作,从 i0 是属于从根到最底下叶子的过程,进行查询累加前缀和操作,从 i前缀和数组.length 属于从底下到根的过程,进行更新操作,和普通实现方法一样,比当前元素大的数都需要递增 1。

树状数组的逻辑是怎么样呢?

首先一个前缀和数组,把它想象成一个二叉树, sums[i] 的值是由底下左右两个节点的和,以此类推一直往下,直到叶子节点,前缀和是 0。叶子节点就是前缀和数组的最后的数。

其次更新只对父级元素负责。因为父级元素的右侧比它小的数就是当前元素。

最后查询 i 的前缀和则只需要通过左右子节点一直累加就行了。

使用 lowbit 位运算能进行二分只是数和数之间有一种数学关系而已,通俗点就是有规律。如果想了解,可以去了解为什么可以成立。如果只是使用,知道他是通过 lowbit 计算二分的树状数组就好。

完整代码如下

    function countSmaller(nums) {
        // 去重,排序,升序是因为方便生成桶
        const a = [...new Set(nums)].sort((a,b) => a - b)
        // 建立映射
        const map = a.reduce((a,b,i) => {
            a[b] = i + 1
            return a
        }, {})
        
        // 生成桶,同时也是前缀和数组
        const b = new Array(a.length).fill(0)
        const result = Array(nums.length)
        for(let i = 0; i < nums.length; i++) {
            // 当前的数映射的下标
            const index = map[nums[i]]
            // 前一个的前缀和就是比当前元素小的个数
            result[i] = getSum(index - 1, b)
            // 更新,比当前元素大的前缀和递增 1
            updateSum(index, b)
        }
        
        return result
    }
    
    function lowbit(n) {
        // 位运算,得出的是 n 的二进制从右边起第一个 1 的位置
        // 比如 6(110) 得出 2
        return n & (n * -1)
    }
    // 树状数组计算前缀和的函数, 从 根n 到 叶子0
    function getSum(n, b) {
        let res = 0
        while(n) {
            res += b[n]
            n -= lowbit(n)
        }
    }
    
    // 树状数组的更新,父级, 从 n 到 length
    function updateSum(n, b) {
        while(n < b.length) {
            b[n]++
            n += lowbit(n)
        }
    }