LeetCode 327. 区间和的个数

题目链接：https://leetcode-cn.com/problems/count-of-range-sum/

题意：给定一个数组nums，找出子数组中元素和在[lower, upper]中有多少个，
[-2, 5, -1]的子数组有[-2, 5, -1]，[-2, 5]，[-2]，[5, -1]，[5]，[-1]共6个，满足条件的共3个

思路：
1.维护一个平衡二叉树，将所有以下标以0为起点的子数组和插入二叉树中
（[-2, 5, -1] = 2，[-2, 5] = 3，[-2] = -2）
满足在[-2, 2]区间内的有2个
2.以下标1为起点的子数组[5, -1], [5]在区间[-2, 2]等价于[-2, 5, -1]，[-2, 5]在区间[-4, 0]

class Program_327 {
    class SplayTree {
        public SplayNode Root;
        public class SplayNode {
            //结点为根的子树中总共有多少个结点
            public int Size;
            //当前结点有多少个Value
            public int Count;
            //当前结点的值
            public long Value;
            //父结点
            public SplayNode Pre;
            //左右子结点
            public SplayNode[] Childs = new SplayNode[2];
            public SplayNode(long _value, SplayNode _pre, SplayNode left = null, SplayNode right = null) {
                Size = 1;
                Count = 1;
                Value = _value;
                Pre = _pre;
                Childs[0] = left;
                Childs[1] = right;
            }
        }
        public SplayTree() {
            //创建两个边界结点
            Root = new SplayNode(long.MinValue, null);
            Root.Childs[1] = new SplayNode(long.MaxValue, Root);
            PushUp(Root);
        }

        /// <summary>
        /// 返回子节点方向
        /// </summary>
        /// <param name="rw"></param>
        /// <returns>rw是父结点的左儿子返回0，右儿子返回1</returns>
        public int F(SplayNode rw) {
            if (rw.Pre.Childs[0] == rw) {
                return 0;
            }
            return 1;
        }

        /// <summary>
        /// 向上更新状态
        /// </summary>
        public void PushUp(SplayNode rt) {
            rt.Size = rt.Count;
            for (int i = 0; i < 2; i++) {
                if (null != rt.Childs[i]) {
                    rt.Childs[i].Pre = rt;
                    rt.Size += rt.Childs[i].Size;
                }
            }
        }

        //将rw旋转到父结点的位置
        public void Rotate(SplayNode rw) {
            int chd = F(rw);
            SplayNode preNode = rw.Pre;

            rw.Pre = preNode.Pre;
            if (null != preNode.Pre) {
                preNode.Pre.Childs[F(preNode)] = rw;
            }
            preNode.Childs[chd] = rw.Childs[1 - chd];
            rw.Childs[1 - chd] = preNode;

            PushUp(preNode);
            PushUp(rw);
        }


        /// <summary>
        /// 调整SplayTree, 将rw调整到rt的子结点
        /// </summary>
        /// <returns></returns>
        public SplayNode Splay(SplayNode rt, SplayNode rw) {
            while (rw.Pre != rt) {
                SplayNode preNode = rw.Pre;
                if (rt == preNode.Pre || F(rw) != F(preNode)) {
                    Rotate(rw);
                } else {
                    Rotate(preNode);
                    Rotate(rw);
                }
            }
            if (null != rt) {
                PushUp(rt);
            } else {
                Root = rw;
            }
            return rw;
        }

        /// <summary>
        /// 插入结点
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public SplayNode Insert(long value) {
            SplayNode curNode = Root;
            SplayNode newNode = new SplayNode(value, null);
            while (true) {
                curNode.Size++;
                int nextChd = 0;
                if (value == curNode.Value) {
                    newNode = curNode;
                    newNode.Count++;
                    break;
                } else if (value < curNode.Value) {
                    nextChd = 0;
                } else {
                    nextChd = 1;
                }
                if (null == curNode.Childs[nextChd]) {
                    newNode.Pre = curNode;
                    curNode.Childs[nextChd] = newNode;
                    break;
                }
                curNode = curNode.Childs[nextChd];
            }
            return newNode;
        }


        /// <summary>
        /// 删除结点
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public bool Delete(long value) {
            SplayNode node = FindRange(value, value);
            if (null == node) {
                return false;
            }
            node.Count--;
            node.Size--;
            if (node.Count <= 0) {
                //删除结点
                node.Pre.Childs[F(node)] = null;
            }
            //减小size
            SplayNode curNode = node;
            while (null != curNode.Pre) {
                curNode = curNode.Pre;
                PushUp(curNode);
            }
            return true;
        }

        /// <summary>
        /// 返回小于value的最大的Node
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public SplayNode FindNode1(long value) {
            SplayNode retNode = null;
            SplayNode curNode = Root;
            while (null != curNode) {
                if (value <= curNode.Value) {
                    curNode = curNode.Childs[0];
                }else {
                    if (null != curNode && (null == retNode || retNode.Value < curNode.Value)) {
                        retNode = curNode;
                    }
                    curNode = curNode.Childs[1];
                }
            }
            return retNode;
        }

        /// <summary>
        /// 返回大于value的最小的Node
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public SplayNode FindNode2(long value) {
            SplayNode retNode = null;
            SplayNode curNode = Root;
            while (null != curNode) {
                if (value < curNode.Value) {
                    if (null != curNode && (null == retNode || curNode.Value < retNode.Value)) {
                        retNode = curNode;
                    }
                    curNode = curNode.Childs[0];
                } else {
                    curNode = curNode.Childs[1];
                }
            }
            return retNode;
        }


        /// <summary>
        /// 将区间内所有结点放到一个子树上，并返回子树根结点
        /// </summary>
        /// <param name="min"></param>
        /// <param name="max"></param>
        /// <returns></returns>
        public SplayNode FindRange(long min, long max) {
            SplayNode minNode = FindNode1(min);
            SplayNode maxNode = FindNode2(max);
            //minNode旋转到Root
            Splay(null, minNode);
            //maxNode旋转为Root右子结点
            Splay(minNode, maxNode);
            return maxNode.Childs[0];
        }
    }

    public int CountRangeSum(int[] nums, int lower, int upper) {
        //前缀和
        long[] Sum = new long[10005];
        //伸展树
        SplayTree splay = new SplayTree();
        Sum[0] = 0;
        for (int i = 1; i <= nums.Length; i++) {
            Sum[i] = Sum[i - 1] + nums[i - 1];
            splay.Splay(null, splay.Insert(Sum[i]));
        }
        int result = 0;
        long tempLower = lower;
        long tempUpper = upper;
        for (int i = 0; i < nums.Length; i++) {
            if (0 < i) {
                splay.Delete(Sum[i]);
            }
            SplayTree.SplayNode node = splay.FindRange(tempLower, tempUpper);
            if (null != node) {
                result += node.Size;
            }
            tempLower += nums[i];
            tempUpper += nums[i];
        }
        return result;
    }
}