abc340 E - Mancala 2

E - Mancala 2

ac library の lazy segtree を使った解法

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace atcoder;
using namespace std;
#define REP(i,n) for(int i=0;i<n;i++)
#define endl '\n'
using ll = long long;

using S = ll;
using F = ll;
S op(S a, S b){ return 0; }
S e(){ return 0; }
S mapping(F f, S x){ return f+x; }
F composition(F f, F g){ return f+g; }
F id(){ return 0; }

int main() {
  int n,m; cin >> n >> m;
  vector<ll> a(n);
  vector<int> b(m);
  REP(i,n) cin >> a[i];
  REP(i,m) cin >> b[i];

  lazy_segtree<S,op,e,F,mapping,composition,id> seg(a);
  REP(i,m) {
    ll now = seg.get(b[i]); // 箱 b[i] に現在入っているボールの数を取得
    ll r = now / n; // 何周するか
    ll rem = now % n; // r 周した余り
    seg.set(b[i],0); // 箱 b[i] を空にする
    seg.apply(0,n,r); // r 周分を全体に加算

    // 余り分を加算
    if (b[i]+rem < n) {
      seg.apply(b[i]+1, b[i]+rem+1,1);
    } else {
      seg.apply(b[i]+1,n,1);
      seg.apply(0, (b[i]+rem)%n+1,1);
    }
  }

  REP(i,n) cout << seg.get(i) << " ";
  return 0;
}

コンテストでは fenwick tree で累積和をして解いたが、lazy segtree の方が簡単だった。