yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

F - Sum of Progression / Codeforces Round 920 (Div. 3)

CF1921 Codeforces C++ Div.3 Codeforces Round 920 Diff 1900 DP 平方分割 競技プログラミング

問題

n 個の整数からなる数列 a が与えられる。 また、整数 s, \, d, \, k によって構成されるクエリが q 個与えられる。

各 q 個のクエリについて、 {a}_{s} + 2 \cdot {a}_{s + d} + \cdots + k \cdot {a}_{s + d \cdot (k - 1)} の値を求めよ。

入力

まず最初の1行目に、テストケースの個数を表す整数 t が与えられる。

その後、 t 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

  • 最初の1行目には、整数 n, \, q が与えられる。
  • 次の1行には、 n 個の整数 {a}_{1}, \, {a}_{2}, \, \cdots , \, {a}_{n} が与えられる。
  • 次の q 行のうち i 行目には、 i 番目のクエリを表す整数 s, \, d, \, k が与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  • 1 \leq t \leq {10}^{4}
  • 1 \leq n \leq {10}^{5}
    • すべてのテストケースにおける n の値の合計は {10}^{5} を超えない。
  • 1 \leq q \leq 2 \cdot {10}^{5}
    • すべてのテストケースにおける q の値の合計は 2 \cdot {10}^{5} を超えない。
  • - {10}^{8} \leq {a}_{i} \leq {10}^{8}
  • 1 \leq s, \, d, \, k \leq n
  • s + d \cdot (k - 1) \leq n

出力

各テストケースについて、 q 個のクエリの答えを順番に空白区切りで1行ずつ出力すること。

解法

まず、クエリについて s, \, d, \, k が与えられた際に、定義通りに {a}_{s} + 2 \cdot {a}_{s + d} + \cdots + k \cdot {a}_{s + d \cdot (k - 1)} を計算した場合の計算量について考えます。

このとき、 k 個の項を利用することから、値を導出するには O(k) の計算量が必要になります。

ここで問題の条件から、 k の値の最大値は s, \, d の値によって決定されます。 すなわち、 \begin{align} s + d \cdot (k - 1) \leq n \end{align} が成立するため、この式を整理すると、 \begin{align} k \leq \frac{n - s}{d} + 1 \end{align} が成立することから、定義通りの計算では \displaystyle O \left( \frac{n}{d} \right) の計算量が必要になると言えます。

このことから、すべての d の値が十分大きいときについてはこの計算を行うのみで実行時間制限に間に合いますが、そうでないときについては実行時間制限には間に合わないと言えます。

そこで、 d の値が小さいときについて、計算量を減らしながらこの計算を行う方法について考えます。 ここで、値 \mathrm{sum}_{i, \, j}, \, \mathrm{dp}_{i, \, j} をそれぞれ \begin{align} \mathrm{sum}_{i, \, j} & = {a}_{i} + {a}_{i + j} + {a}_{i + 2j} + \cdots \\ \mathrm{dp}_{i, \, j} & = {a}_{i} + 2 \cdot {a}_{i + j} + 3 \cdot {a}_{i + 2j} + \cdots \end{align} として定義します。 このとき、まず \mathrm{sum}_{i, \, j} については、 \begin{align} \mathrm{sum}_{i, \, j} & = {a}_{i} + {a}_{i + j} + {a}_{i + 2j} + \cdots \\ & = {a}_{i} + \left( {a}_{i + j} + {a}_{(i + j) + j} + \cdots \right) \\ & = {a}_{i} + \mathrm{sum}_{i + j, \, j} \end{align} が成立します。 従って、 i = n, \, \cdots , \, 1 の順に、 j = 1, \, 2, \, \cdots に対して \mathrm{sum}_{i, \, j} の値を求めていくことができます。

また、 \mathrm{dp}_{i, \, j} の値については、 \begin{align} \mathrm{dp}_{i, \, j} & = {a}_{i} + 2 \cdot {a}_{i + j} + 3 \cdot {a}_{i + 2j} + \cdots \\ & = {a}_{i} + \left( {a}_{i + j} + {a}_{i + 2j} + \cdots \right) + \left( {a}_{i + j} + 2 \cdot {a}_{i + 2j} + \cdots \right) \\ & = {a}_{i} + \left( {a}_{i + j} + {a}_{(i + j) + j} + \cdots \right) + \left( {a}_{i + j} + 2 \cdot {a}_{(i + j) + j} + \cdots \right) \\ & = {a}_{i} + \mathrm{sum}_{i + j, \, j} + \mathrm{dp}_{i + j, \, j} \end{align} が成立します。 従って、 \mathrm{sum}_{i, \, j} が求まっている状態で、 i = n, \, \cdots , \, 1 の順に、 j = 1, \, 2, \, \cdots に対して \mathrm{dp}_{i, \, j} の値を求めていくことができます。

この \mathrm{sum}_{i, \, j}, \, \mathrm{dp}_{i, \, j} が計算されている状態にて、クエリが与えられたときについては、 \begin{align} \, & {a}_{s} + 2 \cdot {a}_{s + d} + \cdots + k \cdot {a}_{s + d \cdot (k - 1)} \\ = & \left( {a}_{s} + 2 \cdot {a}_{s + d} + \cdots + k \cdot {a}_{s + d \cdot (k - 1)} + (k + 1) \cdot {a}_{s + d \cdot k} + \cdots \right) - \left( (k + 1) \cdot {a}_{s + d \cdot k} + \cdots \right) \\ = \, & \mathrm{dp}_{s, \, d} - \left( (k + 1) \cdot {a}_{s + d \cdot k} + (k + 2) \cdot {a}_{s + d \cdot (k + 1)} + \cdots \right) \\ = \, & \mathrm{dp}_{s, \, d} - \left\{ \left( k \cdot {a}_{s + d \cdot k} + k \cdot {a}_{s + d \cdot (k + 1)} + \cdots \right) + \left( {a}_{s + d \cdot k} + 2 \cdot {a}_{s + d \cdot (k + 1)} + \cdots \right) \right\} \\ = \, & \mathrm{dp}_{s, \, d} - k \cdot \mathrm{sum}_{s + d \cdot k , \, d} - \mathrm{dp}_{s + d \cdot k, \, d} \end{align} として整理が行えるので、これは O(1) の計算量によって求めることができます。

一方で、 \mathrm{dp}_{i, \, j} の j の値は最大で n - 1 まで必要になるため、すべての \mathrm{dp}_{i, \, j} の値を求めるには O( {n}^{2} ) の計算量がかかってしまい、これは実行時間制限には間に合いません。

以上から、 j の値の最大値を b とおき、与えられたクエリの d の値が b 以上ならば直接計算し、 b 以下ならば \mathrm{dp}, \, \mathrm{sum} の値を用いて計算することを考えます。

このとき、全体での計算量は \displaystyle O \left( q \cdot \frac{n}{b} + bn \right) ということになります。

これを最小化する b の値については、相加平均・相乗平均の大小関係より、 \begin{align} \frac{nq}{b} & = bn \\ \therefore \; b & = \sqrt{q} \end{align} ということになります。 b がこの値のとき、 \begin{align} q \cdot \frac{n}{b} + bn = q \cdot \frac{n}{\sqrt{q}} + \sqrt{q} \cdot n = 2 n \sqrt{q} \end{align} であるので、 n \sqrt{q} \leq 2 \sqrt{5} \cdot {10}^{7} より、これは実行時間制限に間に合うということが言えます。

この問題では q \leq 2 \cdot {10}^{5} より、 \sqrt{q} \leq 2 \sqrt{5} \cdot {10}^{2} であるので、 b = 300 辺りに固定すれば、 q 個のクエリに実行時間制限に間に合いながら答えることができると言えます。

ソースコード

各テストケースに対し、 q 個のクエリの答えを出力する関数 solve() を以下のように実装しました。

int n, q;
ll s, d, k;
ll a[int(1e5 + 5)], ans[int(2e5 + 5)]; // ans[i] は i 番目のクエリの答えを示す

const int b = 300;
ll dp[int(1e5 + 305)][305], sum[int(1e5 + 305)][305];

void solve() {
  // 値の入力（実行時間制限オーバーを防ぐために、 scanf を使用する）
  scanf("%d %d", &n, &q);
  for (int i = 1; i <= n; i++) {
    scanf("%lld", a + i);
  }

  // i = n, ..., 1 の順に sum, dp の値を更新する
  for (int i = n; i >= 1; i--) {
    for (int j = 1; j <= b; j++) {
      sum[i][j] = sum[i + j][j] + a[i]; // sum[i] は sum[i + j] の値を使用して更新する
      dp[i][j] = dp[i + j][j] + sum[i][j]; // dp[i] は dp[i + j] の値を使用して更新する
    }
  }

  // q 個のクエリの答えを順番に計算していく
  for (int i = 1; i <= q; i++) {
    // クエリの入力
    scanf("%lld %lld %lld", &s, &d, &k);

    if (d <= b) {
      // 入力された d が b 以下の場合は dp, sum を使用して計算する
      ans[i] = dp[s][d];
      if (s + k * d <= n) {
        // 次のインデックスが n 以下の場合のみ、 dp, sum の値から引く
        ans[i] -= dp[s + k * d][d] + k * sum[s + k * d][d];
      }
    } else {
      // 入力された d が b より大きい場合は k 個分の値を順番に加えていく
      for (int j = 0; j < k; j++) {
        ans[i] += ll(j + 1) * a[s + j * d];
      }
    }
  }

  // 答えの出力
  for (int i = 1; i <= q; i++) {
    printf("%lld", ans[i]); // 実行時間制限オーバーを防ぐために、 printf を使用する
    if (i == q) {
      printf("\n");
    } else {
      printf(" ");
    }
  }
  return;
}

感想

ABC335 の F問題 を最近扱いましたが、まさかここでも 平方分割 の問題が出てくるとは思いませんでした。

値の整理の方法についてはやや複雑になる部分もありますが、計算量を減らすために工夫する部分についてはほとんど同じですね。