[LeetCode 315] Count of Smaller Numbers After Self (树状数组 Binary Indexed Tree)

315. Count of Smaller Numbers After Self

You are given an integer array nums and you have to return a new counts array. The counts array has the property where counts[i] is the number of smaller elements to the right of nums[i].

Example:

Given nums = [5, 2, 6, 1]

To the right of 5 there are 2 smaller elements (2 and 1).
To the right of 2 there is only 1 smaller element (1).
To the right of 6 there is 1 smaller element (1).
To the right of 1 there is 0 smaller element.
Return the array [2, 1, 1, 0].

---

题解:

题目的朴素算法很容易想到，遍历数组元素，对每个元素查找其后的比它小的元素个数，得到结果。这个算法的复杂度是O(n^2)的。这个复杂度是会超时的。
想要优化复杂度，可以从每个元素对其后元素的遍历入手，这个遍历实际上可以看成是求和，即小于该元素的计为1，否则计为0，求和。既然是区间求和，就可以尝试使用线段树或树状数组。这里使用实现更简单的树状数组。
首先将给定的nums的拷贝，排序，然后使用一个哈希表来记录每个值对应的下标。建立一个树状数组tree，对nums从后往前遍历，每次对tree的（1，i）求和。求和之后，将这个元素加入tree中，将其值修改为1.
这个算法的复杂度是O(nlogn).

---

代码:

#include 
#include 
#include 
#include 
class BIT {
public:
    BIT(int n) {
        arr = new int[n+1];
        memset(arr, 0, (n+1)*sizeof(int));
        len = n+1;
    }
    ~BIT() {delete arr;}

    void update(int i, int val) {
        while (i < len) {
            arr[i] += val;
            i += i & -i;
        }
    }

    int sum(int i) {
        int ret = 0;
        while (i > 0) {
            ret += arr[i];
            i -= i & -i;
        }
        return ret;
    }
private:
    int *arr;
    int len;
};

class Solution {
public:
    vector<int> countSmaller(vector<int>& nums) {
        int n = nums.size();
        int *tmp_array = new int[n];
        for (int i = 0; i < n; i++) {
            tmp_array[i] = nums[i];
        }
        sort(tmp_array, tmp_array+n);
        unordered_map<int, int> htable;
        for (int i = 0; i < n; i++) {
            htable[tmp_array[i]] = i;
        }
        delete tmp_array;

        BIT tree(n);
        vector<int> ret(n);
        for (int i = n-1; i >= 0; i--) {
            ret[i] = tree.sum(htable[nums[i]]);
            tree.update(htable[nums[i]]+1, 1);
        }
        return ret;
    }
};