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

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

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

問題

 n 組の靴下があり、  i 番目の靴下は色  i の靴下2枚からなる。

このとき、色  {a}_{1}, {a}_{2}, \cdots , {a}_{k} の靴下を1枚ずつ無くしてしまったので、残りの  2n - k 枚の靴下を使って、新たに  \displaystyle \left\lfloor \frac{2n - k}{2} \right\rfloor 組のペアを作り直すことにした。

 i と色  j の靴下のペアの奇妙さを  | i - j | として定義するとき、この奇妙さの総和の最小値を求めよ。

なお、  2 n - k の値が奇数のときは、ペアに使用されない靴下が必ず1枚存在する。

入力

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

次の1行に整数  {a}_{1}, {a}_{2}, \cdots , {a}_{k} が与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  •  1 \leq k \leq n \leq 2 \times {10}^{5}
  •  1 \leq {a}_{1} \lt {a}_{2} \lt \cdots \lt {a}_{k} \leq n

出力

答えを1行で出力すること。

解法

まず、どのような組み合わせが最小値を与えるかについて考えます。

例として、靴下が4枚あるとして、その色が  {a}_{1} \lt {a}_{2} = {a}_{3} \lt {a}_{4} であるとします。

このとき、  ({a}_{1}, {a}_{2}), ({a}_{3}, {a}_{4}) の組み合わせが最善であると考えられ、そのときの奇妙さの総和は \begin{align} | {a}_{1} - {a}_{2} | + | {a}_{3} - {a}_{4} | = ({a}_{2} - {a}_{1} ) + ({a}_{4} - {a}_{3}) \end{align} となります。 ここで、  {a}_{2} = {a}_{3} であるので、 \begin{align} ({a}_{2} - {a}_{1} ) + ({a}_{4} - {a}_{3}) = {a}_{4} - {a}_{1} = | {a}_{1} - {a}_{4} | \end{align} が成立します。 この値は、  ({a}_{1}, {a}_{4}), ({a}_{2}, {a}_{3}) の組み合わせにしたときの奇妙さの総和と等しいです。

このことから、同じ色の組み合わせを作ることは、奇妙さの総和の値に影響を及ぼさないと言えます。 従って、色の揃わない  k 枚の靴下を、奇妙さの総和を最小化するように選んでいくことを考えます。 (このあたりの具体的な証明は公式解説 *1 に詳しく記載されていました。)

さて、  k が偶数の場合は、  ({a}_{1}, {a}_{2}), ({a}_{3}, {a}_{4}), \cdots , ({a}_{k - 1}, {a}_{k}) という組に分けるときに得られる \begin{align} | {a}_{1} - {a}_{2} | + | {a}_{3} - {a}_{4} | + \cdots + | {a}_{k - 1} - {a}_{k} | \end{align} が答えとなります。 そのため、ここからは  k が奇数のときについて考えます。

ここで、  \mathrm{dp}_{i} を、「  {a}_{i} をペアに含めないときの奇妙さの総和」として定めます。 このとき、 \begin{gather} \mathrm{dp}_{1} = | {a}_{2} - {a}_{3} | + | {a}_{4} - {a}_{5} | + \cdots + | {a}_{k - 1} - {a}_{k} | \\ \mathrm{dp}_{2} = | {a}_{1} - {a}_{3} | + | {a}_{4} - {a}_{5} | + \cdots + | {a}_{k - 1} - {a}_{k} | \\ \mathrm{dp}_{3} = | {a}_{1} - {a}_{2} | + | {a}_{4} - {a}_{5} | + \cdots + | {a}_{k - 1} - {a}_{k} | \\ \vdots \\ \mathrm{dp}_{k} = | {a}_{2} - {a}_{3} | + | {a}_{4} - {a}_{5} | + \cdots + | {a}_{k - 2} - {a}_{k - 1} | \end{gather} がそれぞれ成立します。 これらの値の最小値が求める値となります。

この値をそれぞれ求めるとき、  O( {k}^{2}) の計算量がかかってしまい、  k \leq n \leq 2 \times {10}^{5} のため、実行時間制限をオーバーしてしまいます。 従って、既に求められている  \mathrm{dp}_{i} の値を用いて、次の  \mathrm{dp}_{i + 1} の値を求めることを考えます。

ここで、  j j \geq 1 の整数として、  \mathrm{dp}_{2j - 1}, \mathrm{dp}_{2j} の値について着目すると、 \begin{gather} \mathrm{dp}_{2j - 1} = \cdots + |{a}_{2j - 3} - {a}_{2j - 2}| + |{a}_{2j} - {a}_{2j + 1}| + | {a}_{2j + 2} - {a}_{2j +3}| + \cdots \\ \mathrm{dp}_{2j} = \cdots + |{a}_{2j - 3} - {a}_{2j - 2}| + | {a}_{2j - 1} - {a}_{2j + 1}| + | {a}_{2j + 2} - {a}_{2j + 3}| + \cdots \end{gather} がそれぞれ成立するので、 \begin{align} \mathrm{dp}_{2j} = \mathrm{dp}_{2j - 1} - |{a}_{2j} - {a}_{2j + 1}| + |{a}_{2j - 1} - {a}_{2j + 1}| \end{align} が成立します。

同様に、  \mathrm{dp}_{2j}, \mathrm{dp}_{2j + 1} の値について着目すると、 \begin{gather} \mathrm{dp}_{2j} = \cdots + |{a}_{2j - 3} - {a}_{2j - 2}| + | {a}_{2j - 1} - {a}_{2j + 1}| + | {a}_{2j + 2} - {a}_{2j + 3}| + \cdots \\ \mathrm{dp}_{2j + 1} = \cdots + |{a}_{2j - 3} - {a}_{2j - 2}| + |{a}_{2j} - {a}_{2j + 1}| + | {a}_{2j + 2} - {a}_{2j +3}| + \cdots \end{gather} がそれぞれ成立するので、 \begin{align} \mathrm{dp}_{2j + 1} = \mathrm{dp}_{2j} - |{a}_{2j - 1} - {a}_{2j + 1}| + |{a}_{2j} - {a}_{2j + 1}| \end{align} が成立します。

従って、 \mathrm{dp}_{1}, \cdots, \mathrm{dp}_{i - 1} の値がすべて算出されているとすると、 \begin{align} \mathrm{dp}_{i} = \left\{ \begin{array}{cl} \mathrm{dp}_{i - 1} - |{a}_{i - 1} - {a}_{i + 1}| + |{a}_{i} - {a}_{i + 1}| & ( i \text{ が偶数のとき}) \\ \mathrm{dp}_{i - 1} - |{a}_{i - 2} - {a}_{i}| + |{a}_{i - 1} - {a}_{i + 1}| & ( i \text{ が奇数のとき}) \end{array} \right. \end{align} であると言えます。

これを用いてすべての  \mathrm{dp}_{1}, \mathrm{dp}_{2}, \cdots , \mathrm{dp}_{k} を導出し、その内の最小値が求める答えとなります。

 \mathrm{dp}_{1} を求めるのに  O(k) の計算量がかかり、ここからすべての値を導出するのにも  O(k) の計算量がかかることから、この方法で実行時間制限についても十分間に合うと言えます。

ソースコード

main() 関数の中に、答えを出力する部分を直接実装しました。

int n, k;
ll a[int(2e5 + 5)];
ll dp[int(2e5 + 5)];

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

  if (k % 2 == 0) { // k が偶数のときは、順に奇妙さを足したものが答えとなる
    ll ans = 0; // 答えとして最終的に出力する値 ans
    for (int i = 1; i <= k; i += 2) { // i と i + 1 について見るので、 i += 2 とする
      ans += abs(a[i] - a[i + 1]);
    }
    cout << ans << endl; // 答えの出力
    return 0;
  }

  // k が奇数のときは、まずは dp[1] の値を求める
  for (int i = 2; i <= k; i += 2) {
    dp[1] += abs(a[i] - a[i + 1]);
  }

  ll ans = dp[1]; // 答えとして最終的に出力する値 ans
  for (int j = 2; j <= k; j++) {
    dp[j] = dp[j - 1]; // 一旦引き継いでおく
    if (j % 2 == 0) { // j が偶数のとき
      dp[j] -= abs(a[j] - a[j + 1]);
      dp[j] += abs(a[j - 1] - a[j + 1]);
    } else { // j が奇数のとき
      dp[j] -= abs(a[j - 2] - a[j]);
      dp[j] += abs(a[j - 2] - a[j - 1]);
    }
    ans = min(ans, dp[j]); // ans の値の更新
  }
  cout << ans << endl; // 答えの出力

  return 0;
}

感想

B問題 と同様に、場合分けがあるため少し手間取りました。 細かな部分についても無視せずに考慮することの大切さがわかる問題だと思います。