leetcode 1508. 子数组和排序后的区间和

122 阅读1分钟

c++

class Solution {
public:
    struct Data {
        Data(int i, int j, int sum) : i(i), j(j), sum(sum) {}
        int i, j, sum;
    };
    struct CMP {
        bool operator()(Data &a, const Data &b) {
            return a.sum > b.sum;
        }
    };

    int rangeSum(vector<int>& nums, int n, int left, int right) {
        priority_queue<Data, vector<Data>, CMP> que;
        for (int i = 0; i < n; i++) que.push(Data{i, i, nums[i]});
        int ans = 0, mod = 1e9 + 7;
        for (int i = 1; i <= right; i++) {
            Data tmp = que.top();
            que.pop();
            if (i >= left) ans = (ans + tmp.sum) % mod;
            if (tmp.j < n - 1) que.push(Data{tmp.i, tmp.j + 1, tmp.sum + nums[tmp.j + 1]});
        }
        return ans;
    }
};

js

var rangeSum = function(nums, n, left, right) {
    var que = [];
    for (var i = 0; i < n; i++) que.push({i: i, j: i, sum: nums[i]});
    var ans = 0, mod = 1e9 + 7;
    for (var i = 1; i <= right; i++) {
        que.sort((a, b) => a.sum - b.sum);
        var node = que.shift();
        if (i >= left) ans = (ans + node.sum) % mod;
        if (node.j < n - 1) que.push({i: node.i, j: node.j + 1, sum: (node.sum + nums[node.j + 1]) % mod});
    }
    return ans;
};