Golang and BST solution, O(n*logn)


  • 0
    F

    sum[i] = sum(0...i)
    so for sum[i], we need to find how many sum[j] which satisfies j < i and lower<= sum[i]-sum[j] <= upper
    we can put sum[0...j] into a BST, and find how many sum[j].

    type Node struct{
        Val int
        Count int
        lcount, rcount int
        Left, Right *Node
    }
    
    func insert(root *Node, sum int) *Node {
        if root == nil{
            return &Node{sum, 1, 0, 0, nil ,nil}
        }
        if sum < root.Val{
            root.lcount +=1
            root.Left = insert(root.Left, sum)
        }else if sum > root.Val{
            root.rcount +=1
            root.Right = insert(root.Right, sum)
        }else{//sum == root.Val
            root.Count +=1
        }
        return root
    }
    
    func countEqBigger(root *Node,  sum int, eq bool) int{
        if root == nil{
            return 0
        }
        if root.Val == sum{
            if eq{
                return root.Count + root.rcount
            }else{
                return root.rcount
            }
        }else if root.Val > sum{
            return root.Count + root.rcount + countEqBigger(root.Left, sum, eq)
        }else{//root.Val < sum
            return countEqBigger(root.Right, sum, eq)
        }
    }
    
    
    func countRangeSum(nums []int, lower int, upper int) int {
        sums := make([]int, len(nums))
        for i, v :=range nums{
            if i==0{
                sums[i] = v
            }else{
                sums[i] = sums[i-1]+v
            }
        }
        
        var root *Node = nil
        retCount := 0
        for _,sum := range sums{
            c := countEqBigger(root,  sum-upper, true) - countEqBigger(root, sum-lower, false)
            if sum >= lower && sum <= upper{
                c +=1
            }
            retCount += c
            root = insert(root, sum)
        }
        return retCount
    }
    
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.