Untitled
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