Official

E - 図書館の貸出管理 / Library Loan Management Editorial by kyopro_friends


この問題はクエリ先読み+セグメントツリーで解くことができます。

クエリを先読みし、本の返却と問い合わせをまとめて日付の昇順にソートすることで、問題は次のように読み替えられます。

長さ \(N\) の数列 \(A\) があり、最初、全ての要素は \(0\) である。以下の \(2\) 種類のクエリを処理せよ。

  • \(i\) が与えられる。\(A_i\)\(1\) 増やす
  • \(L,R\) が与えられる。\(A_L+\dots+A_R\) を求める

読み替え後の問題はセグメントツリーや Binary Indexed Tree を用いてクエリあたり \(O(\log N)\) で処理することができます。

よって、 \(O((N+Q)\log(N+Q))\) でこの問題を解くことができます。

実装例 (C++)

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

int main(){
  int n, k, q;
  cin >> n >> k >> q;
  vector<array<int, 5>>events;
  // (d[i], 1, i, 0, 0) と (t[i], 2, l[i], r[i], i)
  for(int i=0; i<n; i++){
    int d;
    cin >> d;
    events.push_back({d, 1, i, 0, 0});
  }
  for(int i=0; i<q; i++){
    int l, r, t;
    cin >> l >> r >> t;
    events.push_back({t, 2, l-1, r, i});
  }

  sort(events.begin(), events.end());
  vector<int>ans(q);
  atcoder::fenwick_tree<int>seg(n);
  for(auto[_, t, l, r, i]: events){
    if(t==1){
      seg.add(l, 1);
    }else{
      int ret = seg.sum(l, r);
      if(ret <= k){
        ans[i]= ret;
      }else{
        ans[i] = -1;
      }
    }
  }
  for(int i=0; i<q; i++) cout << ans[i] << endl;
}

実装例 (Python)

from atcoder.fenwicktree import FenwickTree

N, K, Q = map(int, input().split())
events = []
# (d[i], 1, i, 0, 0) と (t[i], 2, l[i], r[i], i)
for i, d in enumerate(map(int, input().split())):
  events.append((d, 1, i, 0, 0))
for i in range(Q):
  l, r, t = map(int, input().split())
  events.append((t, 2, l-1, r, i))

events.sort()
ans = [0] * Q
seg = FenwickTree(N)
for _, t, l, r, i in events:
  if t == 1:
    seg.add(l, 1)
  else:
    ret = seg.sum(l, r)
    if ret <= K:
      ans[i] = ret
    else:
      ans[i] = -1

for a in ans:
  print(a)

posted:
last update: