Untitled

 avatar
unknown
c_cpp
a year ago
2.1 kB
6
Indexable
class Solution {
public:
    int waysToPartition(vector<int>& nums, int k) {
        vector<long long> prefix;
        vector<long long> suffix;
        for (int i = 0; i < nums.size(); ++i) {
            prefix.push_back(nums[i]);
            suffix.push_back(nums[i]);
        }
        map<long long, vector<int>> occ;
        map<int, vector<int>> valid_occ_prefix;
        map<int, vector<int>> valid_occ_suffix;
        for (int i = 0; i < nums.size(); ++i) {
            occ[nums[i]].push_back(i);
        }
        for (int i = 1; i < prefix.size(); ++i)
            prefix[i] += prefix[i - 1];
        for (int i = suffix.size() - 2; i >= 0; --i)
            suffix[i] += suffix[i + 1];
        int res = 0;
        for (int p = 0; p < nums.size() - 1; ++p) {
            if (prefix[p] == suffix[p + 1]) {
                res++;
            }
            long long next_target = suffix[p + 1] - prefix[p] + k;
            auto it = lower_bound(occ[next_target].begin(), occ[next_target].end(), p + 1);
            if (it != occ[next_target].end()) {
                valid_occ_suffix[next_target].push_back(*it);
            }
            long long prev_target = prefix[p] - suffix[p + 1] + k;
            it = lower_bound(occ[prev_target].begin(), occ[prev_target].end(), p + 1);
            if (it != occ[prev_target].begin()) {
                --it;
                valid_occ_prefix[prev_target].push_back(*it);
            }
        }
        for (auto p: occ) {
            for (auto v: p.second) {
                int suff_idx = lower_bound(valid_occ_suffix[p.first].begin(), valid_occ_suffix[p.first].end(), v + 1) - valid_occ_suffix[p.first].begin();
                int suff_cnt = suff_idx;
                int pref_idx = lower_bound(valid_occ_prefix[p.first].begin(), valid_occ_prefix[p.first].end(), v) - valid_occ_prefix[p.first].begin();
                int pref_cnt = valid_occ_prefix[p.first].size() - pref_idx;
                res = max(res, suff_cnt + pref_cnt);
            }
        }
        return res;
    }
};
Editor is loading...
Leave a Comment