<p>Given an integer array <code>nums</code> and two integers <code>lower</code> and <code>upper</code>, return <em>the number of range sums that lie in</em> <code>[lower, upper]</code> <em>inclusive</em>.</p>

<p>Range sum <code>S(i, j)</code> is defined as the sum of the elements in <code>nums</code> between indices <code>i</code> and <code>j</code> inclusive, where <code>i &lt;= j</code>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [-2,5,-1], lower = -2, upper = 2
<strong>Output:</strong> 3
<strong>Explanation:</strong> The three ranges are: [0,0], [2,2], and [0,2] and their respective sums are: -2, -1, 2.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [0], lower = 0, upper = 0
<strong>Output:</strong> 1
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= nums.length &lt;= 10<sup>5</sup></code></li>
	<li><code>-2<sup>31</sup> &lt;= nums[i] &lt;= 2<sup>31</sup> - 1</code></li>
	<li><code>-10<sup>5</sup> &lt;= lower &lt;= upper &lt;= 10<sup>5</sup></code></li>
	<li>The answer is <strong>guaranteed</strong> to fit in a <strong>32-bit</strong> integer.</li>
</ul>

class SegmentTree:
    def __init__(self, size: int):
        self.size = size
        self.tree = [0] * 2 * self.size

    def update(self, idx: int):
        idx += self.size
        while idx > 0:
            self.tree[idx] += 1
            idx //= 2

    def query(self, left: int, right: int) -> int:
        if left == 0 and right == 2:
            print(left + self.size, right + self.size + 1)
        left += self.size
        right += self.size + 1
        count = 0

        while left < right:
            if left % 2 == 1:
                count += self.tree[left]
                left += 1

            if right % 2 == 1:
                right -= 1
                count += self.tree[right]

            left //= 2
            right //= 2

        return count


class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        prefix_sums = {0, -lower, -upper}
        curr_ps = 0
        for num in nums:
            curr_ps += num
            prefix_sums.add(curr_ps)
            prefix_sums.add(curr_ps - lower)
            prefix_sums.add(curr_ps - upper)

        sorted_ps = sorted(list(prefix_sums))

        sgmnt_tree = SegmentTree(size=len(sorted_ps) + 1)

        ps_idx_map = {v: i for i, v in enumerate(sorted_ps)}

        ans = 0
        curr_ps = 0
        sgmnt_tree.update(ps_idx_map[curr_ps])

        for num in nums:
            curr_ps += num
            left = ps_idx_map[curr_ps - upper]
            right = ps_idx_map[curr_ps - lower]
            ans += sgmnt_tree.query(left, right)

            sgmnt_tree.update(ps_idx_map[curr_ps])

        return ans

68386