Count of Range Sum

Count of Range Sum | LeetCode OJ:

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive. Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (ij), inclusive.

Note: A naive algorithm of $O(n^2)$ is trivial. You MUST do better than that.

Example: Given nums = [-2, 5, -1], lower = -2, upper = 2, Return 3. The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

A good interview question! This one could be solved in different ways with various time complexity.

Query of range sum $S(i, j)$ can be facilitated by a pre-calculated prefix sum of $nums$. Properties of prefix sum array(say, $prefix\_sums$) are as follow:

  1. $prefix\_sums[i]$ equals to $\sum_{j=0}^{i-1} nums[j]$. For special case, $prefix\_sums[0]$ equals to zero.
  2. $S(i, j)$ equals to $prefix\_sums[j + 1] - prefix\_sums[i]$.

Given the fact that $S(i, j)$ could be represented by the prefix sum, the original problem statement could be transformed as follow.

For each index $i$ from $0$ to $n - 1$, find the number of valid index $j$ such that:

  1. $j \gt i$.
  2. $lower + prefix\_sums[i] \leq prefix\_sums[j] \leq upper + prefix\_sums[i]$.

In the following sections, different algorithms would be applied to solve the transfromed problem mentioned above. So, if you ever get stuck, come back here to review the problem definition.

From the problem statement we know that the time complexity of finding all indices $j$ with given $i$ must smaller than $O(n)$, otherwise the overall solution will reach $O(n^2)$ since the scale of outer loop is $n$. Intuitively, $O(1)$ solution to the sub-problem is pretty hard to reach(seems impossible). So we will try to find the $O(\log n)$ solutions.

Binary Search Tree

Time complexity:

  • Best case performance $O(n \log n)$.
  • Worst case performance $O(n^2)$.

Recall the transformed question mentioned above, we can construct a binary search tree to satisfy the range query and insert new prefix sum value to BST at each iteration. Following implementation is self-explained.

class Solution {  
 public:
  struct TreeNode {
    TreeNode *left = nullptr;
    TreeNode *right = nullptr;
    int tree_size = 1;
    long long value = 0;
  };

  TreeNode *InsertToBST(TreeNode *root, long long value) {
    if (root == nullptr) {
      auto node = new TreeNode;
      node->value = value;
      return node;
    }

    if (value < root->value) {
      root->left = InsertToBST(root->left, value);
    } else {
      root->right = InsertToBST(root->right, value);
    }
    ++root->tree_size;
    return root;
  }

  int CountGEQ(TreeNode *root, long long hint) {
    // recursive version.
    // if (root == nullptr) {
    //     return 0;
    // }
    //
    // if (root->value < hint) {
    //     return CountGEQ(root->right, hint);
    // } else if (root->value == hint) {
    //     return root->tree_size - (root->left?
    //     root->left->tree_size : 0);
    // } else {
    //     int right_subtree_size = root->right?
    //     root->right->tree_size : 0;
    //     return right_subtree_size + 1 +
    //     CountGEQ(root->left, hint);
    // }

    // iterative version, reduce 40ms.
    int count = 0;
    while (root != nullptr) {
      if (root->value < hint) {
        root = root->right;
      } else if (root->value == hint) {
        count += root->tree_size;
        count -= root->left ? root->left->tree_size : 0;
        break;
      } else {
        count += root->right ? root->right->tree_size : 0;
        count += 1;
        root = root->left;
      }
    }
    return count;
  }

  int CountRange(TreeNode *root, long long lower_bound,
                 long long upper_bound) {
    return CountGEQ(root, lower_bound) -
           CountGEQ(root, upper_bound);
  }

  vector<long long> BuildPrefixSum(
      const vector<int> &nums) {
    vector<long long> prefix_sums(nums.size() + 1, 0);
    for (int i = 1; i <= nums.size(); ++i) {
      prefix_sums[i] = prefix_sums[i - 1] + nums[i - 1];
    }
    return prefix_sums;
  }

  int countRangeSum(vector<int> &nums, int lower,
                    int upper) {
    if (nums.empty()) {
      return 0;
    }
    const int n = nums.size();

    auto prefix_sums = BuildPrefixSum(nums);

    TreeNode *root = nullptr;
    int count = 0;
    for (int i = n; i >= 1; --i) {
      root = InsertToBST(root, prefix_sums[i]);
      count += CountRange(root, prefix_sums[i - 1] + lower,
                          prefix_sums[i - 1] + upper + 1);
    }
    return count;
  }
};

Although this strategy works, but the runtime complexity of BST is highly dependent on the property of dataset, since the cost of searching a BST depends on the height of BST. To solve this problem, we might need to replace BST with some balanced tree, such as red-black tree, but as you know, the cost of implementing a balanced tree is pretty high, especially at the interview.

As a conclusion, using BST to solve the problem works in Leetcode(accepted with 160ms), but might be challenged for its time complexity at the interview.

multiset For Range Search

Time complexity: $O(n \log n + k)$, where $k$ equals to the number of valid ranges. And hence, in the worst case this algorithm has $O(n^2)$ complexity, since there are at most $n(n - 1) / 2$ valid ranges.

multiset is a sorted container with $O(\log n)$ complexity on insertion or accession. Although following code is pretty simple, make sure you understand the difference between lower_bound and upper_bound.

Similar to BST, the worst case performance is still $O(n^2)$, but the actual performance of this solution is not bad in Leetcode(100ms). So I guess the testcases of this question should be enhanced.

class Solution {  
 public:
  vector<long long> BuildPrefixSum(
      const vector<int>& nums) {
    vector<long long> prefix_sums(nums.size() + 1, 0);
    for (int i = 1; i <= nums.size(); ++i) {
      prefix_sums[i] = prefix_sums[i - 1] + nums[i - 1];
    }
    return prefix_sums;
  }

  int countRangeSum(vector<int>& nums, int lower,
                    int upper) {
    if (nums.empty()) {
      return 0;
    }
    const int n = nums.size();

    auto prefix_sums = BuildPrefixSum(nums);

    int count = 0;
    multiset<long long> ordered_sums;
    for (int i = n; i >= 1; --i) {
      ordered_sums.insert(prefix_sums[i]);

      auto lb_iter = ordered_sums.lower_bound(
          lower + prefix_sums[i - 1]);
      auto ub_iter = ordered_sums.upper_bound(
          upper + prefix_sums[i - 1]);

      count += distance(lb_iter, ub_iter);
    }
    return count;
  }
};

Binary Indexed Tree + Binary Search

Time complexity: $O(n \log n)$ in worst case.

If you don't familiar with Binary Indexed Tree, please read following materials first:

The idea is to use BIT to maintain a counter for each prefix sum and update the counter in each iteration, plus a binary search for range query. Since the cost of add/sum operations of BIT require $O(\log n)$ complexity, the worst case time complexity of this solution is guaranteed to be $O(n \log n)$.

Implementation with comments:

class Solution {  
 public:
  // return last index <= hint.
  int BinarySearch(const vector<long long> &ranges,
                   const long long hint) {
    int begin = 0;
    int end = ranges.size();
    while (begin < end) {
      int mid = begin + (end - begin) / 2;
      if (hint < ranges[mid]) {
        end = mid;
      } else {
        begin = mid + 1;
      }
    }
    return end - 1;
  }

  long long SumBIT(const vector<long long> &bit, int idx) {
    long long sum = 0;
    while (idx > 0) {
      sum += bit[idx];
      idx -= (idx & -idx);
    }
    return sum;
  }

  void AddBIT(vector<long long> &bit, int idx, int value) {
    const int n = bit.size();
    while (idx < n) {
      bit[idx] += value;
      idx += (idx & -idx);
    }
  }

  vector<long long> BuildPrefixSum(
      const vector<int> &nums) {
    vector<long long> prefix_sums(nums.size() + 1, 0);
    for (int i = 1; i <= nums.size(); ++i) {
      prefix_sums[i] = prefix_sums[i - 1] + nums[i - 1];
    }
    return prefix_sums;
  }

  vector<long long> BuildSortedRanges(
      const vector<long long> &prefix_sums) {
    // remove duplicates.
    set<long long> help(prefix_sums.cbegin(),
                        prefix_sums.cend());
    // make sure BinarySearch not return zero. This is
    // *important*!
    help.insert(LLONG_MIN);
    // return sorted array.
    return vector<long long>(help.begin(), help.end());
  }

  int countRangeSum(vector<int> &nums, int lower,
                    int upper) {
    if (nums.empty()) {
      return 0;
    }
    const int n = nums.size();

    auto prefix_sums = BuildPrefixSum(nums);
    auto ranges = BuildSortedRanges(prefix_sums);

    vector<long long> bit(ranges.size(), 0);
    for (int i = 0; i <= n; ++i) {
      // mark prefix sums.
      AddBIT(bit, BinarySearch(ranges, prefix_sums[i]), 1);
    }

    int count = 0;
    for (int i = 0; i < n; ++i) {
      // make prefix_sums[i] invisible.
      AddBIT(bit, BinarySearch(ranges, prefix_sums[i]), -1);
      // count the number of valid ranges.
      count += SumBIT(
          bit,
          BinarySearch(ranges, upper + prefix_sums[i]));
      count -= SumBIT(
          bit,
          BinarySearch(ranges, lower + prefix_sums[i] - 1));
    }
    return count;
  }
};

Divide And Conquer

Time complexity: $O(n \log n)$ in worst case.

The strategy is similar to merge sort. Given a interval $[begin, end]$,

  • Divide: we divide the interval into two halves,
  • Conquer: and count the number of valid ranges in each half,
  • Combine: then calculate the number of valid ranges crossing two halves.

Thanks for the sorted property of sub-problem($[begin, mid]$ and $[mid + 1, end]$), the Combine step could be solved in $O(n)$ time complexity using "two pointers" technique. As a consequence, the time complexity of the solution equals to merge sort, which is $O(n \log n)$.

class Solution {  
 public:
  int CountAndSort(vector<long long> &prefix_sums,
                   const int begin, const int end,
                   const int lower, const int upper) {
    // stop condition.
    if (begin == end) {
      return (lower <= prefix_sums[begin]) &&
             (prefix_sums[begin] <= upper);
    }

    // divide and conquer.
    int mid = begin + (end - begin) / 2;
    int count = CountAndSort(prefix_sums, begin, mid, lower,
                             upper) +
                CountAndSort(prefix_sums, mid + 1, end,
                             lower, upper);

    // promise: [begin, mid] and [mid + 1, end] are sorted.
    // count the crossing ranges.
    int j = mid + 1;
    int k = mid + 1;
    for (int i = begin; i <= mid; ++i) {
      while (j <= end &&
             prefix_sums[j] < prefix_sums[i] + lower) {
        ++j;
      }
      while (k <= end &&
             prefix_sums[k] <= prefix_sums[i] + upper) {
        ++k;
      }
      count += k - j;
    }

    // maintain promise.
    inplace_merge(prefix_sums.begin() + begin,
                  prefix_sums.begin() + mid + 1,
                  prefix_sums.begin() + end + 1);
    return count;
  }

  vector<long long> BuildPrefixSum(
      const vector<int> &nums) {
    vector<long long> prefix_sums(nums.size() + 1, 0);
    for (int i = 1; i <= nums.size(); ++i) {
      prefix_sums[i] = prefix_sums[i - 1] + nums[i - 1];
    }
    return prefix_sums;
  }

  int countRangeSum(vector<int> &nums, int lower,
                    int upper) {
    if (nums.empty()) {
      return 0;
    }
    const int n = nums.size();

    auto prefix_sums = BuildPrefixSum(nums);
    return CountAndSort(prefix_sums, 1, n, lower, upper);
  }
};