ABC 371 D - 1D Countryを理解し忘れないようにする。

ABC 371 D - 1D Country

https://atcoder.jp/contests/abc371/tasks/abc371_d

問題概要

数直線上に N 個の村がある。i 番目の村は座標 Xiにあり、Pi人の村人がいます。
Q 個のクエリに答えてください。
i 番目のクエリは以下の形式です。
整数 Li,R iが与えられる。座標が Li以上 Ri以下の村に住んでいる村人の人数の総数を求めよ。

方針

まとめると以下の数式で答えを求められる。
$座標がR以下の場所に住んでいる人の合計 - 座標がL未満の場所に住んでいる人の合計$

ハマった点

• Xiを2分探索するためにソートする必要があるが、それをするとPiとズレが生じてしまう
  • 制約を見れば、すでにXが昇順になっていることが分かる。今後も対応関係のある複数の配列で片方ソートする必要がある場合は、まず制約を確認する。この問題のようにすでに昇順になっているかもしれない。
• 累積和のサイズがN+1なので、配列の添字の指定がやや複雑だった。

解説見た

• 座標XがR以下の合計人数を求める
  • upper_boundを用いることで、累積和のサイズは$|X|+1$なので、結果的に座標がRi以下に住んでいる人の合計人数が求められる。
    • ex: 座標一覧Xが{1,3,5,7}のとき
      • Riが0であれば0番目のイテレータ(0より大きいうち1が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで座標0以下に住んでいる人の合計人数を求められる。累積和の仕様上、0番目のindexを指定すると0になる。
      • Riが1であれば1番目のイテレータ(1より大きいうち3が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって座標1以下に住んでいる人の合計人数を求められる。
      • Riが2であれば1番目のイテレータ(2より大きいうち3が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標1以下に住んでいる人の合計人数を求められる。
      • Riが3であれば2番目のイテレータ(3より大きいうち5が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標3以下に住んでいる人の合計人数を求められる。
      • Riが4の時はRiが3と同じ動作
      • Riが5であれば3番目のイテレータ(5より大きいうち7が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標5以下に住んでいる人の合計人数を求められる。
      • Riが6の時はRiが5と同じ動作
      • Riが7であれば4番目のイテレータ(7より大きいものがないため、一番右側となる)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標7以下に住んでいる人の合計人数を求められる。
      • Riが8であれば4番目のイテレータ(8より大きいものがないため、一番右側となる)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標7以下に住んでいる人の合計人数を求められる。

rep(i,q) {
    int j = upper_bound(all(x),r[i]) - x.begin();
    sum[j]; 
  }

• 座標XがL未満(L-1)の場合の合計人数を求める
  • upper_boundを用いることで、累積和のサイズは$|X|+1$なので、結果的に座標がRi以下に住んでいる人の合計人数が求められる。
    • ex: 座標一覧Xが{1,3,5,7}のとき
      • Liが1のとき、1未満を求めるためにkeyを0とする。そうすると、0番目のイテレータ(0より大きいうち1が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって座標0以下に住んでいる人の合計人数を求められる。座標0は制約上存在しないが、累積和のサイズにより配列外参照にはならない。
      • Liが2のとき、2未満を求めるためにkeyを1とする。そうすると、1番目のイテレータ(1より大きいうち3が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって座標1以下に住んでいる人の合計人数を求められる。
      • Liが3のとき、3未満を求めるためにkeyを2とする。そうすると、1番目のイテレータ(2より大きいうち3が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって、座標2以下に住んでいる人の合計人数を求められる。
      • Liが4のとき、4未満を求めるためにkeyを3とする。そうすると、2番目のイテレータ(3より大きいうち5が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって座標1以下に住んでいる人の合計人数を求められる。
      • Liが5のとき、5未満を求めるためにkeyを4とする。そうすると、3番目のイテレータ(4より大きいうち5が一番左側のため)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって座標5未満、つまり座標3以下に住んでいる人の合計人数を求められる。
    • Liが10のとき、10未満を求めるためにkeyを9とする。そうすると、4番目のイテレータ(8より大きいものがないため、一番右側となる)を取得し、そのイテレータのindex番号を累積和sumに指定することで、累積和とXの配列のサイズの違いによって配列外参照にはならず、座標7以下に住んでいる人の合計人数を求められる。

rep(i,q) {
   int j = upper_bound(all(x),l[i]-1) - x.begin();
   sum[j]; 
 }

提出コード

https://atcoder.jp/contests/abc371/submissions/58600828

int main()
{
  int n;
  cin >> n;
  vector<int> x(n),p(n);
  rep(i,n) cin >> x[i];
  rep(i,n) cin >> p[i];
  int q;
  cin >> q;
  vector<int> l(q),r(q);
  rep(i,q) cin >> l[i] >> r[i];
  vector<ll> sum(n+1);
  rep(i,n) sum[i+1] = sum[i] + p[i];
  rep(i,q) {
    int j1 = upper_bound(all(x),l[i]-1) - x.begin();
    int j2 = upper_bound(all(x),r[i]) - x.begin();
    ll ans = sum[j2]- sum[j1];
    cout << ans << endl;
  }
  return 0;
}