Untitled

#include <bits/stdc++.h>

using namespace std;
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_idx_prefix;
        map<int, vector<int>> valid_occ_idx_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) {
            // prefix == suffix
            if (prefix[p] == suffix[p + 1]) {
                res++;
            }
            // prefix == suffix - arr[i] + k    // arr[i] == suffix - prefix + k
            long long next_target = suffix[p + 1] - prefix[p] + k;
            int idx_occ_target = lower_bound(occ[next_target].begin(), occ[next_target].end(), p + 1) - occ[next_target].begin();  // index p is execluded
            if (idx_occ_target != occ[next_target].size()) {
                // update all occurences of next_target (arr[some_i]) from idx_occ_target to occ[next_target].size() (for all some_i == arr[some_i] from p + 1 to n)
                valid_occ_idx_suffix[next_target].push_back(idx_occ_target);
            }
            // prefix - arr[i] + k == suffix    // arr[i] == prefix - suffix + k
            long long prev_target = prefix[p] - suffix[p + 1] + k;
            idx_occ_target = lower_bound(occ[prev_target].begin(), occ[prev_target].end(), p + 1) - occ[prev_target].begin();  // index p is included
            if (idx_occ_target > 0) {
                // update all occurences of prev_target (arr[some_i]) from 0 to idx_occ_target (for all some_i == arr[some_i] from 0 to p)
                idx_occ_target--;  // index p is included
                valid_occ_idx_prefix[prev_target].push_back(idx_occ_target);
            }
        }
        // update all occurences
        for (auto p: occ) {  // unique numbers
            for (int i = 0; i < p.second.size(); ++i) {
                // count how many suffix
                int suff_idx = lower_bound(valid_occ_idx_suffix[p.first].begin(), valid_occ_idx_suffix[p.first].end(), i + 1) - valid_occ_idx_suffix[p.first].begin();
                int suff_cnt = suff_idx;
                int pref_idx = lower_bound(valid_occ_idx_prefix[p.first].begin(), valid_occ_idx_prefix[p.first].end(), i) - valid_occ_idx_prefix[p.first].begin();
                int pref_cnt = valid_occ_idx_prefix[p.first].size() - pref_idx;
                res = max(res, suff_cnt + pref_cnt);
            }
        }
        return res;
    }
};