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

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

F - Christmas Present 2 / ユニークビジョンプログラミングコンテスト2023 クリスマス (AtCoder Beginner Contest 334)

問題

 xy 平面で表される町があり、サンタと  1 から  n までの番号をつけられた  n 人の子供が住んでいる。 サンタの家の座標は  ({s}_{x}, {s}_{y}) であり、子供  i の家の座標は  {x}_{i}, {y}_{i} である。

サンタは子供の番号順にプレゼントを1個ずつ配る。 子供  i にプレゼントを配るには、プレゼントを1個以上持った状態で子供  i の家に行く必要がある。

しかし、サンタは同時に  k 個までしか同時にプレゼントを持てず、プレゼントを補充するには一旦家に戻る必要がある。

サンタが自分の家を出発し、すべての子供にプレゼントを配り切った状態で再び家に戻るために必要な移動距離の最小値を求めよ。

入力

まず最初の1行目に整数  n, k が与えられる。

次の1行に、整数  {s}_{x}, {s}_{y} が与えられる。

その次の  n 行に渡り、  i 行目には整数  {x}_{i}, {y}_{i} が与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  •  1 \leq k \leq n \leq 2 \times {10}^{5}
  •  - {10}^{9} \leq {s}_{x}, {s}_{y}, {x}_{i}, {y}_{i} \leq {10}^{9}
  •  ({s}_{x}, \, {s}_{y}) \ne ({x}_{i}, \, {y}_{i})
  •  ({x}_{i}, \, {y}_{i}) \ne ({x}_{j}, \, {y}_{j}) (ただし、  i \ne j

出力

移動距離の最小値を1行で出力すること。 ただし、真の値との絶対誤差または相対誤差が  {10}^{-6} 以下のとき正解と判定される。

解法

まず、座標   (x, \, y) (x', \, y') の距離は \begin{align} \sqrt{{(x - x')}^{2} + {(y - y')}^{2}} \end{align} で与えられます。 そのため、サンタの持つことのできるプレゼントの数に制限がない場合、移動距離の最小値は \begin{align} & \sqrt{{({s}_{x} - {x}_{1})}^{2} + {({s}_{y} - {y}_{1})}^{2}} \\ & \quad + \left( \sqrt{{( {x}_{1} - {x}_{2})}^{2} + {( {y}_{1} - {y}_{2})}^{2}} \, + \cdots + \sqrt{{( {x}_{n - 1} - {x}_{n})}^{2} + {( {y}_{n - 1} - {y}_{n})}^{2}} \right) \\ & \quad + \sqrt{{({x}_{n} - {s}_{x})}^{2} + {({y}_{n} - {s}_{y})}^{2}} \\ = & \, \sqrt{{({s}_{x} - {x}_{1})}^{2} + {({s}_{y} - {y}_{1})}^{2}} + \sqrt{{({x}_{n} - {s}_{x})}^{2} + {({y}_{n} - {s}_{y})}^{2}} \\ & \quad + \sum_{i = 2}^{n} \sqrt{{( {x}_{i - 1} - {x}_{i})}^{2} + {( {y}_{i - 1} - {y}_{i})}^{2}} \end{align} によって与えられます。

従って、この値をもとに「  k 個までしかプレゼントを持てない」という条件を加えて考えます。 この条件が加わることにより、総移動距離に対して「一旦自宅を経由する」というペナルティが追加で発生します。

ここで、子供  i - 1 の家から子供  i の家に行く途中に、一旦自宅を経由するときのペナルティを  {p}_{i} と定めると、 \begin{align} {p}_{i} = & \sqrt{{({x}_{i - 1} - {s}_{x})}^{2} + {({y}_{i - 1} - {s}_{y})}^{2}} + \sqrt{{({s}_{x} - {x}_{i})}^{2} + {({s}_{y} - {y}_{i})}^{2}} \\ & \quad - \sqrt{{( {x}_{i - 1} - {x}_{i})}^{2} + {( {y}_{i - 1} - {y}_{i})}^{2}} \end{align} によって与えられます。

さらに、  \mathrm{dp}_{i} を、「条件を満たしながら子供  i の家に到着し、子供  i - 1 の家から子供  i の家に行く途中に一旦自宅を経由したときの、ペナルティの総計の最小値」として定めます。 (ただし、子供  1 の家には必ず条件を満たして到着するため、  \mathrm{dp}_{1} = 0 とします。)

このとき、条件を満たすように移動するには、子供  i - 1 の家に到着した時点で  1 個以上のプレゼントを持っている必要があります。 プレゼントは  k 個まで持てることから、最小で子供  i - k の家に到着する直前に自宅を経由していれば、条件を満たすと言えます。

従って、  \mathrm{dp}_{i} の値については、 \begin{align} \mathrm{dp}_{i} = \min{\left( \mathrm{dp}_{i - k}, \mathrm{dp}_{i - k + 1}, \cdots , \mathrm{dp}_{i - 1} \right)} + {p}_{i} \end{align} によって与えられます。 これを計算していくことで、  \mathrm{dp}_{1}, \cdots , \mathrm{dp}_{n} が得られます。

ここで、最終的には子供  n の家に到着した時点で  1 個以上のプレゼントがある状態になれば、サンタの移動の条件を満たすことになります。 すなわち、子供  n - k + 1 の家に到着する直前以降に1回でも自宅を経由していれば、条件を満たす移動方法ということになります。

従って、最終的なペナルティとして考えられる値は  \mathrm{dp}_{n - k + 1}, \mathrm{dp}_{n - k + 2}, \cdots , \mathrm{dp}_{n} k - 1 個の値であり、このうちの最小値が最終的なペナルティの最小値となります。 すなわち、 \begin{align} & \sqrt{{({s}_{x} - {x}_{1})}^{2} + {({s}_{y} - {y}_{1})}^{2}} + \sqrt{{({x}_{n} - {s}_{x})}^{2} + {({y}_{n} - {s}_{y})}^{2}} \\ & \quad + \sum_{i = 2}^{n} \sqrt{{( {x}_{i - 1} - {x}_{i})}^{2} + {( {y}_{i - 1} - {y}_{i})}^{2}} \\ & \quad + \min{\left( \mathrm{dp}_{n - k + 1}, \mathrm{dp}_{n - k + 2}, \cdots , \mathrm{dp}_{n} \right)} \end{align} がこの問題の答えとして導出される値になります。

ここで、   \min{\left( \mathrm{dp}_{i - k}, \mathrm{dp}_{i - k + 1}, \cdots , \mathrm{dp}_{i - 1} \right)} の値を求める方法について考えます。 これら  k 個の値を毎回見ることとなると、全体で  O(nk) の計算量がかかってしまい、実行時間制限をオーバーしてしまいます。

この問題では子供の家を  i = 1, 2, \cdots , n の順番で訪問していることから、一度使わなくなったインデックスのものはもう二度と使いません。 このことから、C++での priority_queue を用いて管理していくことを考えます。

具体的には、キューの中身を pair<double, int> の型にして、 {dp[i], i} の形で  \mathrm{dp}_{i} の値が計算されるたびに保存していきます。 このような値が保存されているキューに対して、

  1. .top() の2つ目の値が  i - k 未満であれば、その値をキューから取り除く
  2.  i - k 以上のものが存在するまで手順1を繰り返す
  3.  i - k 以上のものが存在したとき、 .top() の1つ目の値が  \mathrm{dp}_{i - k}, \cdots , \mathrm{dp}_{i - 1} での最小値となる
  4.  \mathrm{dp}_{i} の値を計算し、  \{\mathrm{dp}_{i}, i \}pair を新たにキューに追加する

という手順を繰り返すことにより、最終的には  O(n) 程度の計算量で  \mathrm{dp}_{n} までの値を計算できることになります。

これを実装していくことにより、この問題の答えが得られます。

ソースコード

まずは、2点  ({x}_{0}, {y}_{0}),\,({x}_{1}, {y}_{1}) の距離を求める dist() 関数を以下のように実装しました。

double dist(ll x0, ll y0, ll x1, ll y1) {
  return sqrt((x0 - x1) * (x0 - x1) + (y0 - y1) * (y0 - y1)); // 差の2乗の和の平方根を取って返す
}

また、途中で記載した priority_queue を用いて  \mathrm{dp}_{i - k}, \cdots , \mathrm{dp}_{i - 1} の最小値を返す関数 search_min() を以下のように実装しました。

priority_queue<pair<double, int>, vector<pair<double, int>>,
               greater<pair<double, int>>>
    que; // pair<double, int> の値を保存する優先度付きキュー

double search_min(int idx) {
  while (que.top().second < idx - k) {
    que.pop(); // インデックスの値が idx - k 未満である限り、 que から先頭を削除する
  }
  return que.top().first; // que の先頭の1つ目の値を返す
}

以上の関数が実装されている状態で、 main() 関数の中に答えを出力する部分を実装しました。 (答えを出力する際に使用している setprecision については *1 を参照してください。)

int n, k;
ll sx, sy, x[int(2e5 + 5)], y[int(2e5 + 5)];
double ans = 0; // 答えとなる値 ans
double dp[int(2e5 + 5)];

priority_queue<pair<double, int>, vector<pair<double, int>>,
               greater<pair<double, int>>>
    que; // pair<double, int> の値を保存する優先度付きキュー

int main() {
  // 値の入力
  cin >> n >> k;
  cin >> sx >> sy;
  for (int i = 1; i <= n; i++) {
    cin >> x[i] >> y[i];
  }

  // プレゼントの個数が無い状態での距離の総和を求める
  ans += dist(sx, sy, x[1], y[1]);
  for (int i = 2; i <= n; i++) {
    ans += dist(x[i - 1], y[i - 1], x[i], y[i]);
  }
  ans += dist(x[n], y[n], sx, sy);

  que.push({0, 1}); // キューに {0, 1} を追加する
  // i = 2, ... , n に対して、 dp[i] の値を求める
  for (int i = 2; i <= n; i++) {
    double direct = dist(x[i - 1], y[i - 1], x[i], y[i]); // 直接 i - 1 から i へと移った場合の距離
    double via_home =
        dist(x[i - 1], y[i - 1], sx, sy) + dist(sx, sy, x[i], y[i]); // i - 1 から自宅を経由して i へと移った場合の距離

    double min_dist = search_min(i); // search_min() を用いて、dp[i - k], ..., dp[i - 1] からの最小値を求める
    dp[i] = min_dist + via_home - direct; // dp[i] の値を計算する
    que.push({dp[i], i}); // キューに {dp[i], i} の pair を追加する
  }

  // 得られた dp[i] の値から、dp[n - k + 1], ..., dp[n] の最小値を求める
  double min_penalty = 1e18; // ペナルティの値の最小値としての値
  for (int i = max(1, n - k + 1); i <= n; i++) {
    min_penalty = min(min_penalty, dp[i]); // min_penalty を更新する
  }
  ans += min_penalty; // 得られた min_penalty を用いて ans を更新する
  cout << fixed << setprecision(32) << ans << endl; // 小数点以下32桁を表示する

  return 0;
}

感想

公式解説 *2 ではセグメント木を用いた解法が紹介されていますね。 自分なりにセグメント木についても記事を書こうか迷っています。