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 の方が簡単だった。