Simple Heap solution with O(log N) update and query time.


  • 0
    K

    The idea is similar to segment tree, but has a pretty simple implementation.
    Let's take N = 8 as example. We create an array (used as heap) of size N * 2 = 16 that has the following property.

    arr[0] --- not used.
    if p < N, arr[p] = sum of arr[2 * p] and arr[2 * p + 1], and let's call item p is the parent of item p * 2 and item p * 2 + 1.
    if p >= N, arr[p] = nums[p - N].

    so you can see arr[1] = arr[2] + arr[3], arr[3] = arr[6] + arr[7],..., arr[7] = arr[14]+arr[15], arr[8] = nums[0], arr[9] = nums[1], ... arr[15] = nums[7];
    choosing this indexing form has a pretty simple property: for any 1 < p < N-1, the index of its parent is always p / 2.

    class NumArray {
    public:
        int N;
        vector<int> heap;
        NumArray(vector<int> nums) {
            N = nums.size();
            int bit = 1;
            while (N > (1<<bit)) bit++;
            N = (1<<bit);
            heap.resize(N * 2, 0);
            for (int i = 0; i < nums.size(); i++) {
                heap[N+i] = nums[i];
            }
            for (int i = N-1; i >= 1; i--) {
                heap[i] = heap[i *2] + heap[i*2+1];
            }
        }
        
        void update(int i, int val) {
            int diff = val - heap[N+i];
            int p = N + i;
            while (p >= 1) {
                heap[p] += diff;
                p/=2;
            }
        }
        int sum(int len) {
            if (len <= 0) return 0;
            if (len >= N) return heap[1];
            int clen = N, p = 1;
            int s = 0;
            while (clen && len && p < N * 2) {
                if (len == clen) {
                    return s + heap[p];
                } else if (len > clen) {
                    len -= clen;
                    s += heap[p];
                    p = (p+1) * 2;
                    clen /= 2;
                } else {
                    p = p * 2;
                    clen /= 2;
                }
            }
            return s;
        }
        int sumRange(int i, int j) {
            return sum(j+1) - sum(i);
        }
    };
    

Log in to reply
 

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