CanidsSeesaw (その2)

CanidsSeesaw2.cpp

//
using namespace std;
typedef vector<int> vi;
#define sz(a)  int((a).size())
#define pb  push_back
#define rep(var,n)  for(int var=0;var<(n);++var)
#define rep1(var,n)  for(int var=1;var<=(n);++var)
#define all(c)  (c).begin(),(c).end()
#define tr(i,c)  for(auto i=(c).begin(); i!=(c).end(); ++i)
#define found(s,e)  ((s).find(e)!=(s).end())
typedef pair<int,int> ii;

class CanidsSeesaw {
 public:
    vector<int> construct(vector<int> wolf, vector<int> fox, int k) {
        int n = wolf.size();
        vi wa(n); wa[0]=wolf[0]; rep1(i,n-1) wa[i]=wa[i-1]+wolf[i];
        vector<ii> fov(n); rep(i,n) fov[i] = ii(fox[i],i); sort(all(fov));

        int lo = 0, hi = 0;
        int a = 0, ar = 0;
        rep(i, n) {
            a += fov[i].first;
            ar += fov[n-1 - i].first;
            if (a > wa[i]) lo++;
            if (ar > wa[i]) hi++;
        }
        if (k < lo || hi < k) return vector<int>();

        vi fa(n); fa[0]=fov[0].first; rep1(i,n-1) fa[i]=fa[i-1]+fov[i].first;
        int score = lo;
        if (score == k) goto ans;

        for (int t=n-1; t>=0; --t) {
            for (int i=0; i<=t-1; ++i) {
                // i, i+1
                if (fov[i].first < fov[i+1].first) {
                    swap(fov[i], fov[i+1]);
                    bool fwin = fa[i] > wa[i];
                    fa[i] = (i==0 ? 0 : fa[i-1]) + fov[i].first;
                    if (!fwin && fa[i] > wa[i]) ++score;
                    if (score == k) goto ans;
                }
            }
        }
    ans:
        vi fz(n); rep(i,n) fz[i] = fov[i].second; return fz;
    }
};