有序表4:应用：Leetcode第327题区间和的个数

题目链接：力扣 

327. 区间和的个数

难度困难

给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 区间和的个数 。

区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。

示例 1：

输入：nums = [-2,5,-1], lower = -2, upper = 2
输出：3
解释：存在三个区间：[0,0]、[2,2] 和 [0,2] ，对应的区间和分别是：-2 、-1 、2 。

示例 2：

输入：nums = [0], lower = 0, upper = 0
输出：1

提示：

• 1 <= nums.length <= 105
• -231 <= nums[i] <= 231 - 1
• -105 <= lower <= upper <= 105
• 题目数据保证答案是一个 32 位 的整数

解题思路：使用SBTree,但是节点要去重，我们这个题目最重要的方法是查询小于某个key的key的数量，如果不去重的时候会需要一个一个比较，代价较大

下面贴上我写的代码，里面有详细的注释，重点关注下代码中的put add和lessKeysize方法，以及用于去重的set，里面有我自己解题过程中的一些打印和一些和本题本省没有关系的代码，但是对于理解本题有一定的作用，我都留着了

package dataStructure.TreeMap;

import java.util.HashSet;

class SBTreeNode<K extends Comparable> {
    K k;
    SBTreeNode<K> left;
    SBTreeNode<K> right;
    int size;
    int all;
    public SBTreeNode(K k) {
        this.k = k;
        size = 1;
        all = 1;
    }
}

class SBTTreeSet<K extends Comparable> {
    SBTreeNode<K> root;
    HashSet<K> set = new HashSet<>();

    /**
     * 把cur为根的树进行左旋操作：cur的右孩子变成根，
     * cur变成右孩子的左孩子，右孩子的左孩子变成cur右孩子（右孩子如果有左孩子）
     * @param cur
     * @return
     */
    public SBTreeNode<K> leftRotate(SBTreeNode<K> cur) {
        //找到值和cur一样（包含cur）的节点个数
        int same = cur.all - (cur.left == null? 0 : cur.left.all) - (cur.right == null? 0 : cur.right.all);
        //先拿到cur的右孩子，因为下一步即将变更cur的右孩子（变为right的左孩子，如果左孩子存在的话）
        SBTreeNode<K> right = cur.right;
        //cur的右孩子变为right的左孩子，如果左孩子存在的话
        cur.right = right.left;
        right.size = cur.size;
        right.all = cur.all;
        cur.size = (cur.left == null? 0 : cur.left.size) + (cur.right == null? 0 : cur.right.size) + 1;
        cur.all = same + (cur.left == null? 0 : cur.left.all) + (cur.right == null? 0 : cur.right.all);
        //right的左孩子变更为cur（cur要调整成SBTree）
        right.left = cur;
        return right;
    }

    public SBTreeNode<K> rightRotate(SBTreeNode<K> cur) {
        //找到值为cur这个节点的个数
        int same = cur.all - (cur.left == null? 0 : cur.left.all) - (cur.right == null? 0 : cur.right.all);
        //暂存left节点，left节点马上要变更为左子树的右孩子（如果存在的话，不存在就是null）
        SBTreeNode<K> left = cur.left;
        //left节点马上要变更为左子树的右孩子（如果存在的话，不存在就是null）
        cur.left = left.right;
        //left的右孩子指针指向cur
        left.right = cur;
        //把原来以cur为根的子树的size和all信息赋值给left（这些信息原来是cur的size和all存着的）
        left.size = cur.size;
        left.all = cur.all;
        //调整完之后根据cur的left和right来计算cur的新的size和all
        cur.size = (cur.left == null? 0 : cur.left.size) + (cur.right == null? 0 : cur.right.size) + 1;
        cur.all = (cur.left == null? 0 : cur.left.all) + (cur.right == null? 0 : cur.right.all) + same;
        return left;
    }

    //调整函数：根据SB树的定义：当一个节点为根的树的节点数量小于它的侄子节点为根的节点数量的时候，就需要进行平衡型调整
    public SBTreeNode<K> maintain(SBTreeNode<K> cur) {
        int left = cur.left == null? 0 : cur.left.size;
        int right = cur.right == null? 0 : cur.right.size;
        int leftLeft = left == 0? 0 : cur.left.left == null? 0 : cur.left.left.size;
        int leftRight = left == 0? 0 : cur.left.right == null? 0 : cur.left.right.size;
        int rightRight = right == 0? 0 : cur.right.right == null? 0 : cur.right.right.size;
        int rightLeft = right == 0? 0 : cur.right.left == null? 0 : cur.right.left.size;
        if(leftLeft > right) {
            //LL型和LR混合型
            //只有右旋一个操作
            cur = rightRotate(cur);
            //调整完的树和调整之前的只有调整之后的cur和它的右孩子，都分别尝试调整一下
            cur.right = maintain(cur.right);
            cur = maintain(cur);
        } else if(leftRight > right) {
            //纯LR型
            //LR型要先左旋以左孩子为根的树
            cur.left = leftRotate(cur.left);
            //然后右旋整颗以cur为根的树
            cur = rightRotate(cur);
            //这个过程中影响了三个节点（节点的size发生变化），分别是调整之后根cur，左孩子cur.left和右孩子cur.right
            //分别尝试调整
            cur.right = maintain(cur.right);
            cur.left = maintain(cur.left);
            cur = maintain(cur);
        } else if(rightRight > left) {
            //RR型和RL混合型
            //RR型和LL处理方式基本一致（旋转方向相反）
            cur = leftRotate(cur);
            cur.left = maintain(cur.left);
            cur = maintain(cur);
        } else if(rightLeft > left) {
            //纯RL型
            //RL型和LR型处理方式基本一致（旋转方向相反）
            cur.right = rightRotate(cur.right);
            cur = leftRotate(cur);
            cur.left = maintain(cur.left);
            cur.right = maintain(cur.right);
            cur = maintain(cur);
        }
        return cur;
    }

    /**
     * 这里的put不是key, value对，而只是key（真实的要放进去的值）
     * @param k
     * @return
     */
    public SBTreeNode<K> put(K k) {
        //key不能为空，因为要能比较
        if(k == null) {
            throw new RuntimeException("k can't be null");
        }
        //如果这个时候root是空的，那要放进去这个就要变成根节点，当然它在树上是不存在的
        if(root == null) {
            root = add(root, k, false);
        } else {
            //否则先判断存在不存在
            if(set.contains(k)) {
                root = add(root, k, true);
            } else {
                root = add(root, k, false);
            }
        }
        return root;
    }


    /**
     * 在以cur为跟的树上添加以k为key的节点，这个添加过程特殊之处在于
     * 1.如果key已经存在于这个树上，则只更新对应的size(这个key的size要更新，all要更新，沿途的节点更新size)
     * 2.如果key不存在生成这个节点，size为1，all也是1，沿途节点的size和all都加1
     * size是表示这个树上有多少个不同的节点（去重），all表示这个树上一共有多少个节点（包含重复）
     * 整个树上的key都是不重复的
     * @param cur
     * @param k
     * @param contains
     * @return
     */
    public SBTreeNode<K> add(SBTreeNode<K> cur, K k, boolean contains) {
        //这棵树的根为空，直接添加并作为根
        if(cur == null) {
            return new SBTreeNode<>(k);
        }
        //cur是遍历过程中要经过的所有的子树根节点，不管整棵树中是否包含k这个节点，all都要+1
        cur.all ++;
        //不存在的节点才会让size进行+1
        if(!contains) {
            cur.size ++;
        }
        //如果等于当前节点只进行all++的操作，前面我们已经做完++了，这里什么也不用做了
        if(k.compareTo(cur.k) == 0) {
            //cur.all ++;
        } else if(k.compareTo(cur.k) < 0) {
            //如果小于cur的key，到左树上去加，并把根节点返回作为cur的左孩子
            cur.left = add(cur.left, k,contains);
        } else {
            //如果大于cur的key，到右树上去加，并把根节点返回作为cur右孩子
            cur.right = add(cur.right, k, contains);
        }
        //SB树只有在add的时候进行调整平衡性的工作
        cur = maintain(cur);
        return cur;
    }
    //以k为key的节点在SB树中是否存在，这里属于通用操作，目前这个题用不到
    public boolean containsKey(K k) {
        if(root == null || k == null) {
            return false;
        }
        SBTreeNode<K> cur = root;
        boolean contains = false;
        while(cur != null) {
            if(k.compareTo(cur.k) == 0) {
                contains = true;
                break;
            } else if(k.compareTo(cur.k) < 0) {
                cur = cur.left;
            } else {
                cur = cur.right;
            }
        }
        return contains;
    }

    //这里假设节点必然存在,找不到就返回null
    public SBTreeNode<K> getNode(K k) {
        //从root开始
        SBTreeNode<K> cur = root;
        SBTreeNode<K> findNode = null;
        while(cur != null) {
            if(k.compareTo(cur.k) == 0) {
                findNode =  cur;
                break;
            } else if(k.compareTo(cur.k) < 0){
                cur = cur.left;
            } else {
                cur = cur.right;
            }
        }
        return findNode;
    }

    //这个题的解题关键
    //查找比k小的key的数量（总数量，比如有一个节点a比他小，要取a的all，因为原来确实有all个a）
    public int lessKeySize(K k) {
        int lessSize = 0;
        SBTreeNode<K> cur = root;
        //从root开始找比k小的节点
        while(cur != null) {
            //如果刚好等于，那就算一下左子树有多少个节点，加上，结束
            if(k.compareTo(cur.k) == 0) {
                lessSize += cur.left == null? 0 : cur.left.all;
                break;
            } else if(k.compareTo(cur.k) < 0) {
                //如果cur的key比当前要查找的key大，到左子树继续查找，不结算，因为还没找到确定小的空间
                cur = cur.left;
            } else {
                //否则到右树查找，右树查找的时候已经确定当前cur和cur的左边都是比k小的节点
                //都要算到结果集里（cur的all是包含自己、左子树和右子树的所有的节点数，减去右子树的节点数量就是自己和左子树的），然后取cur的右边接续查找
                lessSize += (cur.all - (cur.right == null? 0 : cur.right.all));
                cur = cur.right;
            }
        }
        return lessSize;
    }
}


public class CountOfRangeSum {
    public static int[] generate(int maxSize, int maxValue) {
        int size = (int)(maxSize * Math.random()) + 1;
        int[] arr = new int[size];
        for(int i = 0; i < size; i++) {
            arr[i] = (int)(Math.random() * maxValue) + 1;
        }
        return arr;
    }
    public static void main(String[] args) {
        SBTTreeSet sbtTreeSet = new SBTTreeSet();
        sbtTreeSet.put(1);
        sbtTreeSet.put(1);
        sbtTreeSet.put(2);
        sbtTreeSet.put(3);
        sbtTreeSet.put(3);
        sbtTreeSet.put(3);
        sbtTreeSet.put(5);

        int lessSize5 = sbtTreeSet.lessKeySize(2);
        System.out.println(lessSize5);
        System.out.println(sbtTreeSet.root.all);
        int maxSize = 10;
        int maxValue = 10;
        int[] arr = {1,1,2,3,3,3,5};
        //应该有以下组合:[1,1]   [2]  [1,2] [1 1 2] [2 3] [3] [3] [3] [5]
        long nums = countRangeSum(arr, 2, 5);
        System.out.println(nums);
        int[] arr2 = {2147483647,-2147483648,-1,0};
        //对应的前缀和数组为{2147483647, -1, -2, -2}
        long nums2 = countRangeSum(arr2, -1, 0);
        System.out.println(nums2);
        /*int[] arr = generate(maxSize, maxValue);
        int upper = 8;
        int lower = 2;*/
        //for(int i =)
        System.out.println(Integer.MAX_VALUE);
        System.out.println(Integer.MIN_VALUE);
    }

    public static long countRangeSum(int[] nums, int lower, int upper) {
        //这里一定要使用long，不然的话可能出现越界的情况(元素最大是Integer.MAX_VALUE，最小是Integer.MIN_VALUE，但是再加或者再减可能就越界了)
        long[] sumArr = new long[nums.length];
        SBTTreeSet<Long> sbtTreeSet = new SBTTreeSet();
        //这里一定要先加一个0，因为后面要计算的时候需要考虑之前的累加和，而我们的累加和数组是没有考虑0以前的
        //实际上0以前的累加和就是0
        sbtTreeSet.put(0L);
        int ans = 0;
        for(int i = 0; i < nums.length; i++) {
            sumArr[i] = i == 0? nums[i] : sumArr[i-1] + nums[i];
            //如果累加和是20，lower是5， upper是10，那我们就是要求[20-10,20-5]这个区间的key的个数
            //也就是小于16的key的数量-小于10的key的数量
            long rangeStart = sumArr[i] - upper;
            long rangeEnd = sumArr[i] - lower;
            ans += (sbtTreeSet.lessKeySize(rangeEnd+1) - sbtTreeSet.lessKeySize(rangeStart ));
            //把从0到i的累加和放入SBT中，用于下个i的计算
            sbtTreeSet.put(sumArr[i]);
        }
        return ans;
    }
}

欢迎交流、批量指正