CanidsSeesaw (その1)

CanidsSeesaw1.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 {
    int n;
    vi wa;
    vi fo;
    int p(vi& v){
        int a = 0, c = 0;
        rep(i,n){
            a += fo[v[i]];
            if (a > wa[i]) ++c;
        }
        return c;
    }
    vi middle(vi& lo, vi& hi, int ix=0){
        if (identical(lo, hi)){
            vi mi = lo;
            return mi;
        }
        vi mi; mi.reserve(n);
        set<int> u; rep(x,n) u.insert(x);
        for (int k=0; k<ix; ++k) {
            assert(lo[k] == hi[k]);
            int x = lo[k];
            mi.pb(x);
            u.erase(x);
        }
        for (int k=ix; k<n; ++k){
            if (hi[k] < lo[k]) {
                assert(false);
            }
            if (lo[k] == hi[k]) {
                int x = lo[k];
                mi.pb(x);
                u.erase(x);
                continue;
            }
            // lo[k] < hi[k]
            int m = (lo[k] + hi[k]) / 2;
            auto it = u.lower_bound(m);
            if (m < *it) --it;
            int x = *it;

            if (x < lo[k]) {
                assert(false);
            }
            if (lo[k] < x) {
                assert(x < hi[k]);
                mi.pb(x);
                u.erase(it);

                for (int x : u) mi.pb(x); // fill
                break;
            }
            // x == lo[k] < hi[k]
            vi lo1 = lo;
            if (!next_permutation(lo1.begin()+k+1, lo1.end())) {
                // it was last one
                x = *(++it);
                 assert(lo[k] < x && x <= hi[k]);

                mi.pb(x);
                u.erase(it);
                for (int x : u) mi.pb(x); // fill

                return mi;
            }

            mi.pb(x);
            u.erase(it);

            vi hi9 = mi;
            for (auto it=u.rbegin(); it!=u.rend(); ++it) {
                int x = *it;
                hi9.pb(x);
            }
            return middle(lo1, hi9, k+1);
        }
        assert(mi.size() == n);
        return mi;
    }

    bool identical(vi& a, vi& b){
        if (a.size() != n || b.size() != n) return false;
        rep(i,n) if (a[i] != b[i]) return false;
        return true;
    }

    vector<ii> fov;

    vi bak(vi& mi) {
        vi miv(n);
        rep(i,n) miv[i] = fov[mi[i]].second;
        return miv;
    }
 public:
    vector<int> construct(vector<int> wolf, vector<int> fox, int k) {
        n = wolf.size(); // k = _k;
        wa.resize(n); wa[0]=wolf[0];
        rep1(i,n-1) wa[i]=wa[i-1]+wolf[i];
        fo = fox;

        fov.resize(n);
        rep(i,n) fov[i] = ii(fox[i],i);
        sort(all(fov));

        sort(all(fo));
        vi lo(n), hi(n);
        rep(i,n) { lo[i]=i; hi[i]=n-1-i; }

        int p_lo = p(lo), p_hi = p(hi);
        if (k < p_lo && p_hi < k) return vector<int>();
        if (p_lo == k) return bak(lo);
        if (p_hi == k) return bak(hi);
        // lo < k < hi
        vi last_mi;
        rep(i, 2560) {
            if (identical(lo, hi)) break;
            vi mi = middle(lo, hi);
            if (identical(mi, last_mi)) break;
            int p_mi = p(mi);
            if (p_mi == k) {
                return bak(mi);
            } else if (p_mi < k) {
                if (identical(lo, mi)) break;
                lo = mi;
            } else if (k < p_mi) {
                if (identical(hi, mi)) break;
                hi = mi;
            }
            last_mi = mi;
        }
        return vector<int>();
    }
};