Search code examples
algorithmrangesegment-treebinary-indexed-tree

Find number of items with weight k in a range (with updates and queries)


I am trying to solve the following problem:

Given an array of items with integer weights (arbitrary order), we can have 2 possible operations:

Query: Output the number of items that are of weight k, in the range x to y.
Update: Change the weight of an item at a certain index to v.

Example:

Given the array: [1,2,3,2,5,6,7,3]

If we query for the number of items with weight 2 from index 1 to 3, then the answer would be 2.

If we modify the element at index 2 to have a weight of 2, then we make the same query again, the answer would be 3.

This is certainly a segment tree problem (using point updates). However, I am facing a problem here - each segment will only hold the answer for 1 index. Hence, it seems that I must use vectors in my segment tree. But this would overcomplicate things. Furthermore, I am not sure how to do that either.

Is anyone able to advise me of a better solution?


Solution

  • Instead of segment tree, you should use binary search tree (BST), like AVL, Treap, Splay and etc.

    1. At first, store all the indexes of all appeared values in separated BSTs. In your example [1,2,3,2,5,6,7,3], there should be six BSTs:

      BST 1: [0],
      BST 2: [1,3],
      BST 3: [2,7],
      BST 5: [4],
      BST 6: [5],
      BST 7: [6]

    2. For each query (x, y, k), count the number of elements that fall in range [x, y] in SBT k.

    3. For each update weight[x] = v, remove x from BST weight[x] and add x to BST v

    Time complexity: O(nlogn + mlogn) where n is the length of the data and m is the number of operations.

    Space complexity: O(n)