Untitled
unknown
c_cpp
a year ago
2.1 kB
11
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