LeetCode 327

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
    
        int n = nums.length;
        // 前缀数组
        long[] f = new long[n + 1];
        // 离散化
        // TreeSet:去重+排序
        Set<Long> set = new TreeSet<>();
        set.add(0L);
        for (int i = 1; i <= n; i++) {
            f[i] = f[i - 1] + nums[i - 1];
            set.add(f[i]);
            set.add(f[i] - upper - 1);
            set.add(f[i] - lower);
        }
        // 使用HashMap当作离散化，保存index
        Map<Long, Integer> map = new HashMap<>();
        // 从1开始存
        int cnt = 1;
        for (Long t: set) {
            map.put(t, cnt);
            cnt++;
        }
        BIT tree = new BIT(cnt);
        int res = 0;
        tree.update(map.get(0L), 1);
        for (int i = 1; i <= n; i++) {
            int left = map.get(f[i] - upper - 1), right = map.get(f[i] - lower);
            res += tree.query(right) - tree.query(left);
            tree.update(map.get(f[i]), 1);
        }
        return res;
    }
}

// 树状数组
class BIT {
	int[] tree;
	int n;

	public BIT(int n) {
		this.n = n;
		this.tree = new int[n + 1];
	}

	public int lowbit(int x) {
		return x & -x;
	}

	public void update(int x, int v) {
		for (int i = x; i <= n; i+=lowbit(i))
			tree[i] += v;
	}

	public int query(int x) {
		int res = 0;
		for (int i = x; i > 0; i -= lowbit(i)) 
			res += tree[i];
		return res;
	}
}