D. Petya and Array ～座標圧縮してBITで数える～

• Quiz
  • https://codeforces.com/contest/1042/problem/D
• AC
  • https://codeforces.com/contest/1042/submission/92543146
• 解法
  • editorialの通り、式変形して、leftを固定して条件を満たすrightを数える
  • BITを用いて数え上げるために座標圧縮を使う（累積和を座標圧縮してBITに登録する）
  • leftを1進める前に、今回のleftの分をBITから登録解除する

座標圧縮のはいつもの (table ver)

// 座標圧縮
// 変換テーブルを返す
VI make_compress_table(vector<ll>& A){  
  // 変換表
  auto B = A;
  sort(ALL(B));
  auto it = unique(ALL(B));
  B.erase(it, B.end());
  return B;
}
ll compress_by_table(VI& T, ll v){
  return lower_bound(ALL(T), v) - T.begin();
}

書き換えver

// 座標圧縮
// Aをそのまま書き換えるver
void compress(VI& A){  
  ll N = A.size();
  auto B = A;
  sort(ALL(B));
  
  auto it = unique(ALL(B));
  B.erase(it, B.end());
 
  rep(i,N){
    A[i] = lower_bound(ALL(B), A[i]) - B.begin();
  }
}

code

struct AccSum{
    vector<ll> Ac;
    ll L;
    AccSum(){}
    AccSum(vector<ll> &A){
        L = A.size();
        Ac.resize(L+1);
        FOR(i, 0, L){
            Ac[i+1] = Ac[i] + A[i];
        }
    }
    // sum of [a, b]
    ll sum(ll a, ll b){
        if(a<0) return -1;
        if(b>L-1) return -1;
        return Ac[b+1] - Ac[a];
    }
};

// できること
// 一点追加
// 範囲和
// ei1333's BIT
struct BIT {
  VI data;

  BIT(ll sz) {
    data.assign(++sz, 0);
  }

  ll sum(ll k) {
    ll ret = 0;
    for(++k; k > 0; k -= k & -k) ret += data[k];
    return (ret);
  }

  // [i, j]
  ll range_sum(ll i, ll j){
    if(i==0){
      return sum(j);
    }else{
      return sum(j) - sum(i-1);    
    }
  }

  void add(ll k, ll x) {
    for(++k; k < SZ(data); k += k & -k) data[k] += x;
  }
};

// 座標圧縮
// 変換テーブルを返す
VI make_compress_table(vector<ll>& A){  
  // 変換表
  auto B = A;
  sort(ALL(B));
  auto it = unique(ALL(B));
  B.erase(it, B.end());
  debug(B);
  return B;
}
ll compress_by_table(VI& T, ll v){
  return lower_bound(ALL(T), v) - T.begin();
}

int main(){
    cin.tie(0);
    ios::sync_with_stdio(false);

    // input
    ll N,T;
    cin>>N>>T;

    VI A(N);
    rep(i, N) cin >> A[i];

    AccSum Acc(A);

    BIT bit(410000);

    // compress
    VI V;
    rep(i,N){
      ll v = Acc.sum(0, i);
      V.push_back(v);
      V.push_back(v+T);
    }

    auto tbl = make_compress_table(V);

    // 累積和をBITに登録
    rep(i,N){
      ll v = Acc.sum(0,i);
      ll idx = compress_by_table(tbl, v);
      bit.add(idx, 1);
    }

    ll ans=0;
    rep(l, N){
      ll v = T;
      if(l!=0){
        v += Acc.sum(0,l-1);
      }
      ll limit = compress_by_table(tbl, v);
      
      if(limit>0){
        ans += bit.sum(limit-1);
      }
      ll idx = compress_by_table(tbl, Acc.sum(0,l));
      bit.add(idx, -1); // 登録から消す
    }
    
    p(ans);
   
    return 0;
}