Short & simple O(n log n)


  • 63
    def countRangeSum(self, nums, lower, upper):
        first = [0]
        for num in nums:
            first.append(first[-1] + num)
        def sort(lo, hi):
            mid = (lo + hi) / 2
            if mid == lo:
                return 0
            count = sort(lo, mid) + sort(mid, hi)
            i = j = mid
            for left in first[lo:mid]:
                while i < hi and first[i] - left <  lower: i += 1
                while j < hi and first[j] - left <= upper: j += 1
                count += j - i
            first[lo:hi] = sorted(first[lo:hi])
            return count
        return sort(0, len(first))
    

    First compute the prefix sums: first[m] is the sum of the first m numbers.
    Then the sum of any subarray nums[i:k] is simply first[k] - first[i].
    So we just need to count those where first[k] - first[i] is in [lower,upper].

    To find those pairs, I use mergesort with embedded counting. The pairs in the left half and the pairs in the right half get counted in the recursive calls. We just need to also count the pairs that use both halves.

    For each left in first[lo:mid] I find all right in first[mid:hi] so that right - left lies in [lower, upper]. Because the halves are sorted, these fitting right values are a subarray first[i:j]. With increasing left we must also increase right, meaning must we leave out first[i] if it's too small and and we must include first[j] if it's small enough.

    Besides the counting, I also need to actually merge the halves for the sorting. I let sorted do that, which uses Timsort and takes linear time to recognize and merge the already sorted halves.


  • 0
    C

    Thanks for sharing your solution.
    Two questions:
    What do you mean by ..."is exactly the left-half firsts working with right,"
    why do we need the following operation?
    first[lo:hi] = sorted(first[lo:hi])


  • 0

    What do you mean by ..."is exactly the left-half firsts working with right,"

    What I said right afterwards, in the same sentence.

    why do we need the following operation? first[lo:hi] = sorted(first[lo:hi])

    Because my "mergecounting" needs the halves to be sorted.


  • 0
    C

    Hi, Stefan, Thank you for delivering help from the other half of the world to solve my hour long struggle. Thanks. BTW, always learn so much from your posts here in leetcode. I am looking forward to see more from you.


  • 0

    Hey, how do you know I'm on the other half? :-)

    I just changed my code and explanation a bit, hopefully it's clearer now.


  • 0
    M

    elegant code, elegant explanation. i wish one day i could do this too.


  • 0
    Y

    I think the complexity is beyond O(nlogn), because there are two loops in sort function.


  • 0

    @ymne06 It's O(n log n). That part with the loops only takes linear time.


  • 1

    It's similar to this, which is O(n):

    j = 0
    for i in range(n):
        while j < n:
            j += 1

  • 0
    Y

    From your code above, the i and j will loop n elements, is it O(n) complexity? By the way, I am not familiar with python language.


  • 4

    Yes, it's O(n), because O(2n) = O(n). Here's Java&Co:

    int j = 0;
    for (int i = 0; i < n; i++)
        while (j < n)
            j++;

  • 0
    S

    Cool! Good to know python 'sorted' uses Timsort which is O(n) in this case...


  • 8
    O

    Excellent work! I made more notes to understand the solution:

    With Divide and Conquer, you first suppose you sovled two halves of array fisrt[i] = sum (0..i) and counted the possible qualified ranges (sum of subarray of [lower, upper]) in each half. Then if you find all the qualified range crossing left half and right half then you complete the answer.

    Note you’ve already sorted first[i] in left half and right half. You pick a left element of first array in left half, call it “left”. And you then pick a right in right half. For each left: Right – left is a range sum across left half and right half. Check it meet [lower, upper] or not. If the range sum is too small? Simply increase right. Too large? Stop, you don’t need to increase right pointer anymore, because first is sorted and increasing right can only make the sum larger.

    Therefore for each left, you

    find min right index (i) that meets right – left >= lower,

    find max right index (j) that meets right – left > upper,

    the j – i is the number of valid range sums. Then you iterate next left and search the according answer.

    A side note is when left increase one element, the right-left can only become smaller, what you only need to do is increase right to find larger range sum so that it is valid. Therefore the two loops only takes linear time.


  • 1
    S

    Hi, I have doubt in the next 3 lines:
    First compute the prefix sums: first[m] is the sum of the first m numbers.
    Then the sum of any subarray nums[i:k] is simply first[k] - first[i].
    So we just need to count those where first[k] - first[i] is in [lower,upper].
    Say nums={1,2,3,4,5}, first={0,1,3,6,10,15}, nums[3:5]=first[5]-first[3]=15-6=9.
    However, nums[3]+nums[4]+nums[5]=12. How are they equal? Am I missing something ?

    Thanks in advance.


  • 0

    @suhasm91 nums[3:5] is just nums[3] and nums[4], not nums[5]. Also, nums[5] doesn't even exist (and nums is a list, not a set).


  • 0
    I

    How did you come up with the idea to sort the prefix arrays? I looked at this problem and was only really able to get the naïve solution..


  • 0

    @IWantToPass Just a month earlier, I had come up with a mergesort solution for another problem, where I guess it's more obvious.


  • 0
    I

    Why is j - i the # of valid range sums? Is it cause if u subtract the # of range sums that are less than lower from the # of sums that are greater than upper then u get the remaining sums that are in the range?

    I feel like in a interview this is a pretty hard thing to come up with lol..


  • 0
    I

    In an interview situation, is it ok to use sorted() to merge instead of writing the merge code yourself?


  • 1

    @IWantToPass I don't know. I'm no interviewer and have very little interviewee experience. But in my opinion it definitely should be ok if it's just a small boring part of a larger problem, like it is here. If I were an interviewer, I'd actually find this better. Everybody knows how to write merging yourself, yawn. But using sorted (and noting its linear runtime) is simpler, probably faster, and shows some creativity/boldness and extra knowledge.


Log in to reply
 

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