CSES - Hotel Queries - O(log n)

#include <bits/stdc++.h>
using namespace std;

#define f first
#define s second
#define pb push_back
#define rep(i, begin, end)                                                     \
  for (__typeof(end) i = (begin) - ((begin) > (end));                          \
       i != (end) - ((begin) > (end)); i += 1 - 2 * ((begin) > (end)))
#define db(x) cout << '>' << #x << ':' << x << endl;
#define sz(x) ((int)(x).size())
#define newl cout << "\n"

#define ll long long int
#define vi vector<int>
#define vll vector<long long>
#define vvll vector<vll>
#define pll pair<long long, long long>

#define fast_io()                                                              \
  ios_base::sync_with_stdio(false);                                            \
  cin.tie(NULL);                                                               \
  cout.tie(NULL);

const int mxN = 2e5 + 5;
const int mxM = 2e5 + 7;
const int lg = 22;
ll N = 0;

template <class T> struct Seg {
  const T ID = 0;
  // comb(ID,b) = b
  T comb(T a, T b) { return max(a, b); }
  vector<int> roots;
  int n;
  vector<T> seg;
  void init(int _n) {
    n = _n;
    seg.assign(2*n, ID);
  }

  void init_roots(){
    vector<int> roots_r;
    for(auto l = n, r = n<<1; l < r; l >>= 1, r >>= 1){
      if(l & 1) roots.push_back(l++);
      if(r & 1) roots_r.push_back(--r);
    }
    roots.insert(roots.end(), roots_r.rbegin(), roots_r.rend());
  }

  void pull(int p) { seg[p] = comb(seg[p<<1], seg[p<<1^1]); }

  void upd(int p, T val) { // set val at position p
    seg[p += n] = val;
    for (p = p>>1; p; p = p>>1) pull(p);
  }

  T qry(int l, int r) { // on interval [l, r]
    T ra = ID, rb = ID;
    for (l += n, r += n + 1; l < r; l = l>>1, r = r>>1) {
      if (l & 1) ra = comb(ra, seg[l++]);
      if (r & 1) rb = comb(seg[--r], rb);
    }
    return comb(ra, rb);
  }

  // log(n) * log(n)
  ll findfirst_1(ll rn) {
    int l = -1, r = n-1;
    while(r-l > 1) {
      auto mid = (r+l)>>1;
      if(qry(0, mid) < rn) l = mid;
      else r = mid;
    }
    if(qry(r, r) >= rn) return r;
    return -1;
  }

  // log(n); need biggest power of 2 in iterative segment tree, greater than equal to n;
  ll findfirst_2(ll x) {
    // return any value from [0, n-1] if present; else return -1;
    ll ans = -1;
    ll i = 1;
    while(i < n && seg[i] >= x) {
      if(seg[i<<1] >= x) i = i<<1;
      else if(seg[i<<1^1] >= x) i = i<<1^1;
      else return -1;
    }
    if(seg[i] >= x) return i-n;
    return -1;
  }

  // log(n), does not need n to be power of 2
  ll findfirst_3(ll x) {
    if(qry(0, n-1) < x) return -1;
    ll j = 0; // iterate every root;
    while(j < seg.size() && seg[roots[j]] < x) {
      j++;
    }
    // found the first root, containing el>=x;
    ll i = roots[j];
    while(i < n && seg[i] >= x) {
      if(seg[i<<1] >= x) i = i<<1;
      else i = i<<1^1;
    }
    return i-n;
  }
};

ll tc, n, m, k;

Seg<ll> st;


int main() {
  fast_io();
#ifndef ONLINE_JUDGE
  freopen("../input.txt", "r", stdin);
  freopen("../output.txt", "w", stdout);
#endif

  cin>>n>>m;

//  for(int i = 0; N = 1<<i, 1<<i < n; i++);
//  st.init(N);

  st.init(n);
  st.init_roots();

  rep(i, 0, n) {
    ll x; cin>>x;
    st.upd(i, x);
  }

  rep(i, 0, m) {
    ll rn; cin>>rn;
    ll ans = 0;
    ans = st.findfirst_3(rn);
    if(ans != -1) {
      cout<<ans+1<<" ";
      st.upd(ans, st.qry(ans, ans)-rn);
    } else {
      cout<<0<<" ";
    }
  }

  return 0;
}