高级树_LeetCode.327_区间和的个数

🌸题目

🍁给定一个整数数组 nums，返回区间和在 [lower, upper] 之间的个数，包含 lower 和 upper。
区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。

说明:
最直观的算法复杂度是 O(n2) ，请在此基础上优化你的算法。

示例:

1
2
3

输入: nums = [-2,5,-1], lower = -2, upper = 2,
输出: 3 
解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2。

🌸解法一：归并排序

设前缀和数组为 $\textit{preSum}$，则问题等价于求所有的下标对(i,j)，满足

preSum[j]−preSum[i]∈[lower,upper]

我们先考虑如下的问题：给定两个升序排列的数组 n1,n2，试找出所有的下标对 (i,j)，满足n2[j]−n1[i]∈[lower,upper]

在已知两个数组均为升序的情况下，这一问题是相对简单的：我们在 n2中维护两个指针 l,r。起初，它们都指向 n 2 的起始位置。随后，我们考察 n 1 的第一个元素。首先，不断地将指针 l 向右移动，直到$$ n_2[l] \ge n_1[0] + \textit{lower}$$为止，此时， ll 及其右边的元素均大于或等于 $n_1[0] + \textit{lower}$；随后，再不断地将指针 r 向右移动，直到 $n_2[r] > n_1[0] + \textit{upper}$为止，则 r左边的元素均小于或等于 $n_1[0] + \textit{upper}$,故区间 [l,r) 中的所有下标 j，都满足n2[j]−n1[0]∈[lower,upper]

接下来，我们考察 n 1 的第二个元素。由于n1 是递增的，不难发现 l,r 只可能向右移动。因此，我们不断地进行上述过程，并对于 n 1中的每一个下标，都记录相应的区间 [l,r) 的大小。最终，我们就统计得到了满足条件的下标对 (i,j) 的数量。

在解决这一问题后，原问题就迎刃而解了：我们采用归并排序的方式，能够得到左右两个数组排序后的形式，以及对应的下标对数量。对于原数组而言，若要找出全部的下标对数量，只需要再额外找出左端点在左侧数组，同时右端点在右侧数组的下标对数量，而这正是我们此前讨论的问题。

代码

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long s = 0;
        long[] sum = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i) {
            s += nums[i];
            sum[i + 1] = s;
        }
        return countRangeSumRecursive(sum, lower, upper, 0, sum.length - 1);
    }

    public int countRangeSumRecursive(long[] sum, int lower, int upper, int left, int right) {
        if (left == right) {
            return 0;
        } else {
            int mid = (left + right) / 2;
            int n1 = countRangeSumRecursive(sum, lower, upper, left, mid);
            int n2 = countRangeSumRecursive(sum, lower, upper, mid + 1, right);
            int ret = n1 + n2;

            // 首先统计下标对的数量
            int i = left;
            int l = mid + 1;
            int r = mid + 1;
            while (i <= mid) {
                while (l <= right && sum[l] - sum[i] < lower) {
                    l++;
                }
                while (r <= right && sum[r] - sum[i] <= upper) {
                    r++;
                }
                ret += r - l;
                i++;
            }

            // 随后合并两个排序数组
            int[] sorted = new int[right - left + 1];
            int p1 = left, p2 = mid + 1;
            int p = 0;
            while (p1 <= mid || p2 <= right) {
                if (p1 > mid) {
                    sorted[p++] = (int) sum[p2++];
                } else if (p2 > right) {
                    sorted[p++] = (int) sum[p1++];
                } else {
                    if (sum[p1] < sum[p2]) {
                        sorted[p++] = (int) sum[p1++];
                    } else {
                        sorted[p++] = (int) sum[p2++];
                    }
                }
            }
            for (int j = 0; j < sorted.length; j++) {
                sum[left + j] = sorted[j];
            }
            return ret;
        }
    }
}

🌸解法二：线段树

依然考虑前缀和数组$ \textit{preSum}$

对于每个下标 j，以 j 为右端点的下标对的数量，就等于数组 $\textit{preSum}$[0..j-1] 中的所有整数，出现在区间$ [\textit{preSum}[j]-\textit{upper}, \textit{preSum}[j]-\textit{lower}]$ 的次数。故很容易想到基于线段树的解法。

我们从左到右扫描前缀和数组。每遇到一个数$ \textit{preSum}[j]$，我们就在线段树中查询区间$ [\textit{preSum}[j]-\textit{upper}, \textit{preSum}[j]-\textit{lower}]$内的整数数量，随后，将 $\textit{preSum}[j]$ 插入到线段树当中。

注意到整数的范围可能很大，故需要利用哈希表将所有可能出现的整数，映射到连续的整数区间内。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long sum = 0;
        long[] preSum = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i) {
            sum += nums[i];
            preSum[i + 1] = sum;
        }
        
        Set<Long> allNumbers = new TreeSet<Long>();
        for (long x : preSum) {
            allNumbers.add(x);
            allNumbers.add(x - lower);
            allNumbers.add(x - upper);
        }
        // 利用哈希表进行离散化
        Map<Long, Integer> values = new HashMap<Long, Integer>();
        int idx = 0;
        for (long x : allNumbers) {
            values.put(x, idx);
            idx++;
        }

        SegNode root = build(0, values.size() - 1);
        int ret = 0;
        for (long x : preSum) {
            int left = values.get(x - upper), right = values.get(x - lower);
            ret += count(root, left, right);
            insert(root, values.get(x));
        }
        return ret;
    }

    public SegNode build(int left, int right) {
        SegNode node = new SegNode(left, right);
        if (left == right) {
            return node;
        }
        int mid = (left + right) / 2;
        node.lchild = build(left, mid);
        node.rchild = build(mid + 1, right);
        return node;
    }

    public int count(SegNode root, int left, int right) {
        if (left > root.hi || right < root.lo) {
            return 0;
        }
        if (left <= root.lo && root.hi <= right) {
            return root.add;
        }
        return count(root.lchild, left, right) + count(root.rchild, left, right);
    }

    public void insert(SegNode root, int val) {
        root.add++;
        if (root.lo == root.hi) {
            return;
        }
        int mid = (root.lo + root.hi) / 2;
        if (val <= mid) {
            insert(root.lchild, val);
        } else {
            insert(root.rchild, val);
        }
    }
}

class SegNode {
    int lo, hi, add;
    SegNode lchild, rchild;

    public SegNode(int left, int right) {
        lo = left;
        hi = right;
        add = 0;
        lchild = null;
        rchild = null;
    }
}

🌸解法三：动态增加节点的线段树

与方法二类似，但我们可以不实用哈希表进行映射，而是只在线段树的插入操作过程中动态地增加树中的节点。而当我们进行查询操作时，如果到达一个空节点，那么说明对应的区间中暂时还没有值，就可以直接返回 0。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long sum = 0;
        long[] preSum = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i) {
            sum += nums[i];
            preSum[i + 1] = sum;
        }
        
        long lbound = Long.MAX_VALUE, rbound = Long.MIN_VALUE;
        for (long x : preSum) {
            lbound = Math.min(Math.min(lbound, x), Math.min(x - lower, x - upper));
            rbound = Math.max(Math.max(rbound, x), Math.max(x - lower, x - upper));
        }
        
        SegNode root = new SegNode(lbound, rbound);
        int ret = 0;
        for (long x : preSum) {
            ret += count(root, x - upper, x - lower);
            insert(root, x);
        }
        return ret;
    }

    public int count(SegNode root, long left, long right) {
        if (root == null) {
            return 0;
        }
        if (left > root.hi || right < root.lo) {
            return 0;
        }
        if (left <= root.lo && root.hi <= right) {
            return root.add;
        }
        return count(root.lchild, left, right) + count(root.rchild, left, right);
    }

    public void insert(SegNode root, long val) {
        root.add++;
        if (root.lo == root.hi) {
            return;
        }
        long mid = (root.lo + root.hi) >> 1;
        if (val <= mid) {
            if (root.lchild == null) {
                root.lchild = new SegNode(root.lo, mid);
            }
            insert(root.lchild, val);
        } else {
            if (root.rchild == null) {
                root.rchild = new SegNode(mid + 1, root.hi);
            }
            insert(root.rchild, val);
        }
    }
}

class SegNode {
    long lo, hi;
    int add;
    SegNode lchild, rchild;

    public SegNode(long left, long right) {
        lo = left;
        hi = right;
        add = 0;
        lchild = null;
        rchild = null;
    }
}

🌸解法四：树状数组

树状数组与线段树基于类似的思想，不过树状数组支持的基本查询为求出 $[0, \textit{val}]$ 之间的整数数量。为了查询区间$ [\textit{preSum}[j]-\textit{upper}, \textit{preSum}[j]-\textit{lower}]$内的整数数量，需要执行两次查询，即分别查询 $[0, \textit{preSum}[j]-\textit{upper}-1]$ 区间的整数数量 L 和$[0,\textit{preSum}[j]-\textit{lower}]$ 区间的整数数量 R，答案即为两者作差 R-L。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long sum = 0;
        long[] preSum = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i) {
            sum += nums[i];
            preSum[i + 1] = sum;
        }
        
        Set<Long> allNumbers = new TreeSet<Long>();
        for (long x : preSum) {
            allNumbers.add(x);
            allNumbers.add(x - lower);
            allNumbers.add(x - upper);
        }
        // 利用哈希表进行离散化
        Map<Long, Integer> values = new HashMap<Long, Integer>();
        int idx = 0;
        for (long x: allNumbers) {
            values.put(x, idx);
            idx++;
        }

        int ret = 0;
        BIT bit = new BIT(values.size());
        for (int i = 0; i < preSum.length; i++) {
            int left = values.get(preSum[i] - upper), right = values.get(preSum[i] - lower);
            ret += bit.query(right + 1) - bit.query(left);
            bit.update(values.get(preSum[i]) + 1, 1);
        }
        return ret;
    }
}

class BIT {
    int[] tree;
    int n;

    public BIT(int n) {
        this.n = n;
        this.tree = new int[n + 1];
    }

    public static int lowbit(int x) {
        return x & (-x);
    }

    public void update(int x, int d) {
        while (x <= n) {
            tree[x] += d;
            x += lowbit(x);
        }
    }

    public int query(int x) {
        int ans = 0;
        while (x != 0) {
            ans += tree[x];
            x -= lowbit(x);
        }
        return ans;
    }
}

🌸解法五：平衡二叉搜索树

考虑一棵平衡二叉搜索树。若其节点数量为 N，则深度为 O(log N)。二叉搜索树能够在 O(logN) 的时间内，对任意给定的值 $\textit{val}$,查询树中所有小于或等于该值的数量。

因此，我们可以从左到右扫描前缀和数组。对于$ \textit{preSum}[j]$ 而言，首先进行两次查询，得到区间 $[\textit{preSum}[j]-\textit{upper}, \textit{preSum}[j]-\textit{lower}]$内的整数数量；随后再将 $\textit{preSum}[j]$ 插入到平衡树中。

平衡二叉搜索树有多种不同的实现，最经典的为 AVL 树与红黑树。此外，在算法竞赛中，还包括 Treap、SBT 等数据结构。

下面给出基于 Treap 的实现。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long sum = 0;
        long[] preSum = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i) {
            sum += nums[i];
            preSum[i + 1] = sum;
        }
        
        BalancedTree treap = new BalancedTree();
        int ret = 0;
        for (long x : preSum) {
            long numLeft = treap.lowerBound(x - upper);
            int rankLeft = (numLeft == Long.MAX_VALUE ? (int) (treap.getSize() + 1) : treap.rank(numLeft)[0]);
            long numRight = treap.upperBound(x - lower);
            int rankRight = (numRight == Long.MAX_VALUE ? (int) treap.getSize() : treap.rank(numRight)[0] - 1);
            ret += rankRight - rankLeft + 1;
            treap.insert(x);
        }
        return ret;
    }
}

class BalancedTree {
    private class BalancedNode {
        long val;
        long seed;
        int count;
        int size;
        BalancedNode left;
        BalancedNode right;

        BalancedNode(long val, long seed) {
            this.val = val;
            this.seed = seed;
            this.count = 1;
            this.size = 1;
            this.left = null;
            this.right = null;
        }

        BalancedNode leftRotate() {
            int prevSize = size;
            int currSize = (left != null ? left.size : 0) + (right.left != null ? right.left.size : 0) + count;
            BalancedNode root = right;
            right = root.left;
            root.left = this;
            root.size = prevSize;
            size = currSize;
            return root;
        }

        BalancedNode rightRotate() {
            int prevSize = size;
            int currSize = (right != null ? right.size : 0) + (left.right != null ? left.right.size : 0) + count;
            BalancedNode root = left;
            left = root.right;
            root.right = this;
            root.size = prevSize;
            size = currSize;
            return root;
        }
    }

    private BalancedNode root;
    private int size;
    private Random rand;

    public BalancedTree() {
        this.root = null;
        this.size = 0;
        this.rand = new Random();
    }

    public long getSize() {
        return size;
    }

    public void insert(long x) {
        ++size;
        root = insert(root, x);
    }

    public long lowerBound(long x) {
        BalancedNode node = root;
        long ans = Long.MAX_VALUE;
        while (node != null) {
            if (x == node.val) {
                return x;
            }
            if (x < node.val) {
                ans = node.val;
                node = node.left;
            } else {
                node = node.right;
            }
        }
        return ans;
    }

    public long upperBound(long x) {
        BalancedNode node = root;
        long ans = Long.MAX_VALUE;
        while (node != null) {
            if (x < node.val) {
                ans = node.val;
                node = node.left;
            } else {
                node = node.right;
            }
        }
        return ans;
    }

    public int[] rank(long x) {
        BalancedNode node = root;
        int ans = 0;
        while (node != null) {
            if (x < node.val) {
                node = node.left;
            } else {
                ans += (node.left != null ? node.left.size : 0) + node.count;
                if (x == node.val) {
                    return new int[]{ans - node.count + 1, ans};
                }
                node = node.right;
            }
        }
        return new int[]{Integer.MIN_VALUE, Integer.MAX_VALUE};
    }

    private BalancedNode insert(BalancedNode node, long x) {
        if (node == null) {
            return new BalancedNode(x, rand.nextInt());
        }
        ++node.size;
        if (x < node.val) {
            node.left = insert(node.left, x);
            if (node.left.seed > node.seed) {
                node = node.rightRotate();
            }
        } else if (x > node.val) {
            node.right = insert(node.right, x);
            if (node.right.seed > node.seed) {
                node = node.leftRotate();
            }
        } else {
            ++node.count;
        }
        return node;
    }
}

🌸解法一说明：为什么可以归并

看完后我明白如果数组是有序的，那么可以容易的求出区间数量，那么为什么对数组排序前后所求的区间数量不会改变呢，下面我来说明一下（官方并没有实际说明这一点）

初始归并：
此时只会有0或1个元素，不涉及左右两段的情况，是可以的

合并归并：此时是有左右两段的，左右两段是分别有序的，对前缀和数组排序并不会修改数组中元素的值，只是改变了元素是位置，如对left ~ right=3~5位置的前缀和排序，排序后前缀和3 ~ 5位置的数还是原来3 ~ 5位置的数，只是排列变化了

设想一个一般的情况，现在是某一层的递归，左，右两段区间left ~ mid, mid+1 ~ right的符合要求的区间数量已经通过countRangeSumRecursive计算了出来，整个left ~ right区间中可能的符合要求的区间情况是两端点在left ~ mid中；两端点在mid+1 ~ right；一个端点在left~ mid中，一个端点在mid+1~ right中，所以现在只要求出第三种情况的区间数量就可以了
通过上面的说明，left ~ mid,mid+1~ right区间中的数还是原来区间中的数，只是顺序变成了有序，而有序是容易计算符合要求的区间数量的，一个图说明为什么第三种情况排序前后符合数量的区间数量是不变的

🌸指路：如何学习本题的算法和数据结构

忽略官方题解中的方法一，剩余的四种方法分别使用了线段树、树状数组和平衡树。这些方法都不是面试的考点，甚至在笔试中也很少出现，所以大部分读者应该是完全不知道这些都是啥神奇的数据结构的。所以我觉得这里有必要补充以下两点：

这道题需要哪些接口。
上面的这些神奇的数据结构可以提供哪些接口。

题目描述

给定数组 A，它的长度为 n，对应的元素以及下标为A[0],A[1],⋯,A[n−1]。

令 S(i,j) 为A[i] 到 A[j] 的和，即

$S(i, j) = \sum_{k=i}^j A[k]$

题目需要求出满足$ \textit{lower} \leq S(i, j) \leq \textit{upper}$ 的二元组 (i, j)(i,j) 的个数。

换成人话就是，问数组 A 有多少个连续的子数组，其元素只和在 $[\textit{lower}, \textit{upper}]$ 的范围内。

题目分析

暴力的做法是使用前缀和。令 P 为 A 的前缀和数组，那么

S(i, j) = P[j] - P[i-1]

可以在O(1) 的时间求出。这里我们规定边界 P[-1] = 0。

这样一来，我们枚举所有的二元组 (i, j)，算出S(i,j) 并判断其是否在范围内即可。时间复杂度为$ O(n^2)$。

那么怎么进行优化呢？我们考虑从小到大枚举 j，由于

$\textit{lower} \leq P[j] - P[i-1] \leq \textit{upper}$

我们可以得到 P[i-1] 应该满足的不等式

$P[j] - \textit{upper} \leq P[i-1] \leq P[j] - \textit{lower}$

因此本质上，我们需要一个数据结构支持下面的两个操作：

操作 1「查询」：给定一个范围 $[\textit{left}, \textit{right}]$，查询数据结构中在该范围内的元素个数。对应到本题中，我们给定的范围就是$ \big[P[j] - \textit{upper}, P[j] - \textit{lower}\big]$；

操作 2「更新」：给定一个元素 xx，我们需要把它添加到数据结构中。对应到本题中，我们给定的元素就是 P[j]。

如果有了这样一个数据结构，我们就可以很方便地做出本题：

我们首先将 00 放入数据结构中，随后我们从小到大枚举 j，查询 $\big[P[j] - \textit{upper}, P[j] - \textit{lower}\big]$ 范围内的元素个数并计入答案。在查询完成之后，我们将P[j] 添加进数据结构，就可以进行下一次枚举。

频数数组

很多数据结构都是基于「频数数组」。

给定数组 t 以及它的下标范围 [L, R][L,R]，t[x] 就表示元素 x 在数据结构中的出现次数。基于此，上述的两个操作可以变为：

操作 1「查询」：给定一个范围$ [\textit{left}, \textit{right}]$，查询$ t[\textit{left}]$到 $t[\textit{right}]$ 的和；

操作 2「更新」：给定一个元素 x，将 t[x] 增加 1。

这也是线段树和树状数组的基础，它们需要的空间都与数组 t 的下标范围 [L, R][L,R] 正相关。在本题数据规模较大的情况下（例如测试数据中，出现了元素值达到 32 位有符号整数的上下界），线段树和树状数组都会超出空间限制，因此需要借助「离散化」操作，将这些元素映射到一个较小规模的区间内。

离散化

给定数组元素 [1, 22, 333, 4444, 55555]，如果我们只关注它们之间的大小关系，那么该数组其实和 [1, 2, 3, 4, 5] 是等价的。

这就是离散化的技巧。我们将所有会涉及到比较操作的数全部放入一个列表中，进行排序，再从小到大依次给它们赋予一个新的值。在离散化完成后，任何两个数之间的相对大小都不会改变，但是元素的上下界范围被限制住，这使得我们可以方便地使用一些数据结构。

在本题中，我们可以将所有的 $P[j], P[j] - \textit{upper}, P[j] - \textit{lower}$一起进行离散化，并将它们从小到大，依次赋予一个从 1 开始的整数值。这样一来，我们所有涉及到的元素值都会在 [1, 3(n+1)]的范围内。

线段树

当我们将元素离散化后，就可以直接使用线段树了。最基础的线段树恰好就支持这两种操作：

操作 1「查询」：给定一个范围$ [\textit{left}, \textit{right}]$,查询$ t[\textit{left}] $到 $t[\textit{right}]$ 的和；

操作 2「更新」：给定一个元素 x，将 t[x] 增加 $\delta$。

我们只需要时刻令$\delta$=1即可。两种操作的时间复杂度均为$ O(\log n)$

树状数组

当我们将元素离散化后，也可以直接使用树状数组了。最基础的线段树支持这两种操作：

操作 1「查询」：给定一个下标$ \textit{right}$，查询 t[1] 到$ t[\textit{right}]$ 的和（即前缀和）；

操作 2「更新」：给定一个元素 x，将t[x] 增加 $\deltaδ$。

我们只需要时刻令 $\delta=1$ 即可，并且通过调用操作 1 两次（即 $\textit{right}$ 和$ \textit{left}-1left−1）相减得到 t[\textit{left}]$ 到 $t[\textit{right}]$ 的和。两种操作的时间复杂度均为 $O(\log n)$。

平衡树

平衡树实际上就是「平衡」的二叉搜索树，它与线段树和树状数组不同，并且它不需要借助离散化操作。支持的操作（在本题中会使用到的）主要有以下几种：

操作 1「lower bound」：给定一个元素 x，查询平衡树中最小的大于等于 x 的元素；

操作 2「upper bound」：给定一个元素 x，查询平衡树中最小的大于 x 的元素；

操作 3「rank」：给定一个元素 x（它必须在平衡树中），求它是第几小的元素。当存在重复元素时，会计入多次；

操作 4「insert」：给定一个元素 x，将它放入平衡树中。

所有操作的时间复杂度均为 $O(\log n)$。大部分语言自带的平衡树支持操作 1 和 2 和 4 但不支持操作 3。

那么对于本题中需要的两种操作：

「查询」：我们令 $u = P[j] - \textit{upper}$，$v = P[j] - \textit{lower}$。对 u 使用操作 11 得到 u’′
，对 v 使用操作 2 得到 v’ 。我们再使用操作 3 得到 u’的$ rank r_u$ 以及 v’
的$ rank r_v$ ，那么$ r_v’-r_u’$′，就是$ \big[P[j] - \textit{upper}, P[j] - \textit{lower}\big]$ 中的元素个数。

「更新」：我们对 x 使用操作 4 即可。

学习方法

关于这些竞赛难度的知识点，我建议读者在学有余力的情况下学习。这些知识点对面试几乎没有任何帮助；相反，在没有完全掌握这些知识点的前提下，可能会影响读者原本正常的思维，产生「看什么题都是线段树」之类的后果。

这里我推荐两个参考资料：

OI Wiki，是一个信息学竞赛爱好者用爱发电的算法小百科。
ac-library，是日本著名算法竞赛平台 AtCoder 整理的算法模板，但其中没有平衡树。

读者也可以参考其它互联网上的博客。线段树、树状数组和平衡树在算法竞赛圈中是非常基础的知识点，优质的博客数量也很多

🌸题目