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

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

C - Closest Cities / Educational Codeforces Round 161 (Rated for Div. 2)

問題

 n 個の町が数直線上に並んでおり、  i 番目の町は座標  {a}_{i} に位置している。  n 個の町は数直線上に昇順に並んでいるため、  {a}_{1} \lt {a}_{2} \lt \cdots \lt {a}_{n} である。

このとき、2つの町  x, \, y の距離は  | {a}_{x} - {a}_{y} | である。 町  i に対して、最も近い町を  j とするとき、  k \ne i である  k に対して  | {a}_{i} - {a}_{j} | \leq | {a}_{i} - {a}_{k} | が成立する。

ここで、あなたは町  x から町  y に移動しようとしている。 このとき移動に際して、以下の行動を行うことができる。

  • 他の町  y に対して、  | {a}_{x} - {a}_{y} | 枚のコインを支払って移動する。
  •  x に最も近い町に、  1 枚のコインを支払って移動する。

この状況下で、  m 個のクエリが与えられる。 各クエリに対して2つの町  x, \, y が与えられるので、町  x から町  y に移動するのに必要な最小のコインの枚数を求めよ。

入力

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

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

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

条件

  • 実行時間制限: 2s
  • メモリ制限: 256MB
  •  1 \leq t \leq {10}^{4}
  •  2 \leq n \leq {10}^{5}
    • すべてのテストケースにおける  n の値の合計は  {10}^{5} を超えない。
  •  0 \leq {a}_{1} \lt {a}_{2} \lt \cdots \lt {a}_{n} \leq {10}^{9}
  •  1 \leq m \leq {10}^{5}
    • すべてのテストケースにおける  m の値の合計は  {10}^{5} を超えない。
  •  1 \leq {x}_{i}, \, {y}_{i} \leq n
    • ただし、  {x}_{i} \ne {y}_{i} である。

出力

各テストケースに対して、  m 行に渡り、  i 行には  i 番目のクエリでの払うコインの最小枚数を出力すること。

解法

まず、町  x から町  y へと進む際に、  x \lt y であれば「町  x → 町  x + 1 → … → 町  y 」と進むことが最小コストとなりますし、  x \gt y であれば「町  x → 町  x - 1 → … → 町  y 」と進むことが最小コストとなります。

これは例えば  x \lt y の場合、仮に町  x - 1 を経由した場合では、「町  x → 町  x - 1 → 町  x → 町  x + 1 → … → 町  y 」というルートになり、最初の「町  x → 町  x - 1 → 町  x 」の部分が余分なものとなってしまうからです。 ( x \gt y の場合についても同様です。)

ここで、  \mathrm{cost 0}_{i} を町  i + 1 から町  i まで進むときの最小コスト、  \mathrm{cost 1}_{i} を町  i - 1 から町  i まで進むときの最小コストとして定めます。 ただし、  \mathrm{cost 0}_{n} = 0, \, \mathrm{cost 1}_{1} = 0 とします。

まず、  \mathrm{cost 0}_{i} の値について考えます。

 n は必ず町  n - 1 が最も近い町となることから、  \mathrm{cost 0}_{n - 1} = 1 が成立します。

その上で、  1 \leq i \leq n - 2 とするとき、  \mathrm{cost 0}_{i} の値は、町  i + 1 の最も近い町が町  i になる場合には  1 となり、そうでない場合は  | {a}_{i} - {a}_{i + 1} | となります。 すなわち、 \begin{align} \mathrm{cost 0}_{i} = \left\{ \begin{aligned} & 1 & & (|{a}_{i + 1} - {a}_{i}| \leq |{a}_{i + 1} - {a}_{i + 2}| \, \text{のとき}) \\ & | {a}_{i} - {a}_{i + 1} | & & (\text{otherwise}) \end{aligned} \right. \end{align} が成立します。

以上から、  \mathrm{cost 0}_{i} の値が  i = 1, \, 2, \, \cdots , \, n に対して導けます。

次に、  \mathrm{cost 1}_{i} の値について考えます。

 1 は必ず町  2 が最も近い町となることから、  \mathrm{cost 1}_{2} = 1 が成立します。

その上で、  3 \leq i \leq n とするとき、  \mathrm{cost 1}_{i} の値は、町  i - 1 の最も近い町が町  i になる場合には  1 となり、そうでない場合は  |{a}_{i - 1} - {a}_{i}| となります。 すなわち、 \begin{align} \mathrm{cost 1}_{i} = \left\{ \begin{aligned} & 1 & & (|{a}_{i - 1} - {a}_{i}| \leq |{a}_{i - 1} - {a}_{i - 2}| \, \text{のとき}) \\ & | {a}_{i} - {a}_{i - 1} | & & (\text{otherwise}) \end{aligned} \right. \end{align} が成立します。

以上から、  \mathrm{cost 1}_{i} の値が  i = 1, \, 2, \, \cdots, \, n に対して導けます。

このようにして求まった  \mathrm{cost0}, \, \mathrm{cost1} を用いて、町  x から町  y までの最小コストを求めます。

 x \gt y の場合は、町  i + 1 から町  i への移動を繰り返し行うことから、 \begin{align} \mathrm{cost 0}_{x - 1} + \mathrm{cost 0}_{x - 2} + \cdots + \mathrm{cost 0}_{y} \end{align} によって最小コストが求まります。

また、  x \lt y の場合は、町  i - 1 から町  i への移動を繰り返し行うことから、 \begin{align} \mathrm{cost 1}_{x + 1} + \mathrm{cost 1}_{x + 2} + \cdots + \mathrm{cost 1}_{y} \end{align} によって最小コストが求まります。

これらを毎回足していくことで導出すると、  O(n) 程度の計算量が1回につき必要になり、全体で [O(nm)] の計算量が必要になるため、これでは実行時間制限に間に合いません。

従って、値  \mathrm{sum0}_{i}, \, \mathrm{sum1}_{i} を、 \begin{align} \mathrm{sum 0}_{i} & = \mathrm{cost 0}_{n} + \mathrm{cost 0}_{n - 1} + \cdots + \mathrm{cost 0}_{i} \\ \mathrm{sum 1}_{i} & = \mathrm{cost 1}_{1} + \mathrm{cost 1}_{2} + \cdots + \mathrm{cost 1}_{i} \end{align} としてそれぞれ定めます。

この値を用いることによって、  x \gt y の場合の最小コストは、 \begin{align} \mathrm{cost 0}_{x - 1} + \cdots + \mathrm{cost 0}_{y} = \mathrm{sum 0}_{y} - \mathrm{sum 0}_{x} \end{align} で導出することができ、  x \lt y の場合の最小コストは、 \begin{align} \mathrm{cost 1}_{x + 1} + \cdots + \mathrm{cost 1}_{y} = \mathrm{sum 1}_{y} - \mathrm{sum 1}_{x} \end{align} で導出することができます。

 \mathrm{sum 0}_{i}, \, \mathrm{sum 1}_{i} の値はそれぞれ  O(n) で求めることができ、その上で1つのクエリに対しては  O(1) で求めることができるので、この方法では全体で  O(n + m) の計算量で求めることができます。 従って、この方法によって実行時間制限に十分間に合うと言えます。

ソースコード

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

int n, m, a[int(1e5 + 5)];

int cost[int(1e5 + 5)][2], sum[int(1e5 + 5)][2]; // cost[i][0] が cost0 の値を表す。 sum についても同様

void solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }

  // cost0 の値を求める
  for (int i = 1; i <= n - 2; i++) {
    if (abs(a[i + 1] - a[i]) <= abs(a[i + 1] - a[i + 2])) {
      cost[i][0] = 1;
    } else {
      cost[i][0] = abs(a[i] - a[i + 1]);
    }
  }
  cost[n - 1][0] = 1;

  // cost1 の値を求める
  cost[2][1] = 1;
  for (int i = 3; i <= n; i++) {
    if (abs(a[i - 1] - a[i]) <= abs(a[i - 1] - a[i - 2])) {
      cost[i][1] = 1;
    } else {
      cost[i][1] = abs(a[i] - a[i - 1]);
    }
  }

  // 求めた cost0, cost1 の値から、 sum0, sum1 の値を求める
  for (int i = 1; i <= n; i++) {
    sum[i][1] = sum[i - 1][1] + cost[i][1];
  }
  for (int i = n; i >= 1; i--) {
    sum[i][0] = sum[i + 1][0] + cost[i][0];
  }

  // クエリへの回答
  cin >> m;
  for (int i = 0; i < m; i++) {
    int x, y;
    cin >> x >> y;

    if (x <= y) {
      // x <= y のときは、 sum1 の値を用いる
      cout << sum[y][1] - sum[x][1] << endl;
    } else {
      // x > y のときは、 sum0 の値を用いる
      cout << sum[y][0] - sum[x][0] << endl;
    }
  }

  return;
}

感想

それぞれの道順に対して、  i の値がどのときにどうなるのか、ということを考えるのが少し複雑になりました。

記事を書くときのように整理できていると実装もしやすくなるのですが、いざコンテスト本番でこの方針で行うとなると、正解なのかどうかやや不安になりそうです。