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

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

B - 250 Thousand Tons of TNT / Codeforces Round 909 (Div. 3)

問題

 n 個の箱があり、  i 番目の箱は  {a}_{i} トンの重さである。

この  n 個の箱を  k 個ずつ順番に運び、トラックに積むことを考える。 すなわち、

  • 最初の  k 個の箱を最初のトラックに積む。
  • 次の  k 個の箱を2番目のトラックに積む。
  •  \cdots
  • 最後の  k 個の箱を  \frac{n}{k} 番目のトラックに積む。

ということを行う。 ただし、どのトラックも必ず  k 個の箱が積まれているものとする。

すべてのトラックに荷物を等しい個数分けられるようなすべての  k に対して、各トラックに積まれた荷物の重さの最大値と最小値の差の最大値として考えられる値を求めよ。 ただし、トラックが1台のみの場合は、最大値と最小値の差を  0 とする。

入力

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

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

  • 最初の1行目には、整数  n が与えられる。
  • 次の1行には、 n 個の整数  {a}_{i} が順に与えられる。

条件

  •  1 \leq t \leq {10}^{4}
  •  1 \leq n \leq 1.5 \cdot {10}^{5}
    • すべてのテストケースにおける  n の値の合計は  1.5 \cdot {10}^{5} を超えない。
  •  1 \leq {a}_{i} \leq {10}^{9}

出力

 t 行にわたり、各テストケースでの答えを1行ごとに出力すること。

解法

まず、すべてのトラックに等しい個数だけ荷物を積まなければならないので、  k n の約数であるということが言えます。

そのため、整数  n の約数をすべて列挙する必要がありますが、これは  O(\sqrt{n}) 程度の計算量で行えます。 以前書いた記事 にて、ある整数に対する約数を求める方法を実装したので、今回の問題でもこれと同様のものを使用します。

このようにして求まった  k に対して、問題文通りに  ({a}_{1} + \cdots + {a}_{k}),  ({a}_{k + 1} + \cdots + {a}_{2k}),  \cdots ,  ({a}_{n - k + 1} + \cdots + {a}_{n}) を計算することで、その最大値と最小値の差を求めることができます。 これをすべての  k に対して行うことで、差の最大値を求めることができます。

この方法では、 n の約数の個数を  d(n) としたとき、  O(n \cdot d(n)) 程度の計算量で求めることができます。  n \leq 1.5 \cdot {10}^{5} より、 d(n) の値の最大値は  n = 110880 のときの  d(n) = 144 なので、これは実行時間制限に間に合うと考えられます。

ソースコード

解法内でも引用した、 以前書いた記事 と同様に、整数 num の約数を set<int> の形で保存することを考えました。 そのため、ある整数 num の約数をすべて求めて保存するための関数 prepare_divisors() を以下のように実装しました。

set<int> divisors[int(2e5 + 5)]; // 約数を保存するための set
bool visited[int(2e5 + 5)]; // すでに約数を計算していたら true にする

void prepare_divisors(int num) {
  if (visited[num]) {
    return; // すでに約数を計算していたら特に何もせず return する
  }

  visited[num] = true; // num の約数を計算したことにする
  for (int i = 2; i * i <= num; i++) { // i は 2 から √ num までを調べれば OK
    if (num % i != 0) {
      continue; // i が num の約数でないときは次の i を調べる
    }

    prepare_divisors(num / i); // 整数 (num / i) の約数を調べる
    for (int d : divisors[num / i]) {
      // 整数 (num / i) の約数に対して、その整数と、それに i を掛けた値が num の約数となる
      divisors[num].insert(d);
      divisors[num].insert(i * d);
    }
    break;
  }

  // num が素数や1の場合も考慮して、 1 と num を約数に入れておく
  divisors[num].insert(1);
  divisors[num].insert(num);
  return;
}

この prepare_divisors() 関数がある状態で、答えとなる最大値を出力する solve() 関数を以下のように実装しました。

int n;
ll a[int(2e5 + 5)];
set<int> divisors[int(2e5 + 5)]; // 約数を保存するための set

ll solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  prepare_divisors(n); // 整数 n の約数をすべて用意する

  ll ans = 0; // 最終的に答えとなる値
  for (int k : divisors[n]) { // n の約数 k について見る
    int cnt = 0; // cnt は現在のトラックに何個目の箱を積んだかを表す
    // now_weight は現在の重量、min_weight はトラックごとの重量の最小値、 max_weight はトラックごとの重量の最大値をそれぞれ表す
    ll now_weight = 0, min_weight = ll(1e18), max_weight = 0;

    for (int i = 1; i <= n; i++) {
      cnt++;
      now_weight += a[i];

      if (cnt == k) { // トラックに積んだのが k 個目の荷物であれば、以下の値を更新していく
        min_weight = min(min_weight, now_weight); // min_weight の更新
        max_weight = max(max_weight, now_weight); // max_weight の更新
        cnt = 0; // cnt の値をリセットする
        now_weight = 0; // now_weight の値をリセットする
      }
    }

    ans = max(ans, max_weight - min_weight); // 得られた max_weight, min_weight から ans の値を更新する
  }

  return ans; // ans を返す
}

感想

約数の列挙については、別に set を使用しなくても for ループを回していくなどしてできますが、最近 別の問題 で触れたので、このようにしました。