Good Array

 avatar
Saraseruu
c_cpp
a month ago
1.6 kB
2
Indexable
Never
#include <bits/stdc++.h>
#define endl "\n"
#define ll long long
#define IO ios::sync_with_stdio(0); cin.tie(0); cout.tie(0)
using namespace std;

vector<ll> prefix_v(vector <int> v) {   // O(n)
    vector<ll> s(v.size());
    s[0]= v[0];
    for(int i=1; i<v.size(); ++i) {
        s[i] = s[i-1] + v[i];
    }
    return s;
}


int main() {
    IO;
    int n; cin>>n;
    multimap <int,int> mp;  // multimap allows for similar keys
    vector <int> vec, indices;
    int item;
    for(int i=0; i<n; ++i) {
        cin>>item;
        mp.insert(pair<int,int>(item, i+1));
    }
    for(auto pair: mp) {
        //cout<<pair.first<<" : "<<pair.second<<endl;
        vec.push_back(pair.first);
        indices.push_back(pair.second);
    }
    // cout<<endl;
    vector <ll> pre = prefix_v(vec);

    // cout<<"vector is: "<<endl;
    // for(int i: vec) cout<<i<<" ";
    // cout<<endl;
    // cout<<"prefix is: "<<endl;
    // for(int i: indices) cout<<i<<" ";

    int count =0;
    vector <int> result;
    if (pre[n-2]- pre[0] == vec[n-1]) {  // handling i = 0
        ++count;
        result.push_back(indices[0]);
    }
    if (pre[n-3] == vec[n-2]) { //handling i = n-1
        ++count;
        result.push_back(indices[n-1]);
    }
     for(int i=1;i<n-1; ++i) {
         if (pre[n-2]- pre[i] + pre[i-1] == vec[n-1]) {  // 2 cases to be handled 
             ++count;
             result.push_back(indices[i]);
         }
    }
    cout<<count<<endl;
    for(int i: result)  cout<<i<<" ";


    return 0;
}

Leave a Comment