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

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

D - Very Different Array / Codeforces Round 920 (Div. 3)

問題

 n 個の整数からなる数列  a と、  m 個の整数からなる数列  b がある。(ただし、  n \geq m とする。)

数列  b から  n 個の要素を選び、新たに  n 個の整数からなる数列  c を作成する。 このとき数列  a, \, c に対して、  D = \displaystyle \sum_{i = 1}^{n} | {a}_{i} - {c}_{i} | と定めるとき、  D の最大値を求めよ。

入力

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

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

  • まず最初の1行目に、整数  n, \, m が与えられる。
  • 次の1行には、  n 個の整数  {a}_{1}, \, {a}_{2}, \, \cdots , \, {a}_{n} が順に与えられる。
  • 次の1行には、  m 個の整数  {b}_{1}, \, {b}_{2}, \, \cdots , \, {b}_{m} が順に与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 256MB
  •  1 \leq t \leq 100
  •  1 \leq n \leq  m \leq 2 \cdot {10}^{5}
    • すべてのテストケースにおける  m の値の合計は  2 \cdot {10}^{5} を超えない。
  •  1 \leq {a}_{i}, \, {b}_{i} \leq {10}^{9}

出力

各テストケースに対して、  D の最大値を1行で出力すること。

解法

一般性を失わないため、数列  a を昇順にソートされたものとして考えます。 すなわち、  {a}_{1} \leq {a}_{2} \leq \cdots \leq {a}_{n} として考えます。

ここで、 \begin{align} D & = \displaystyle \sum_{i = 1}^{n} | {a}_{i} - {c}_{i} | \\ & = | {a}_{1} - {c}_{1} | + | {a}_{2} - {c}_{2} | + \cdots + | {a}_{n} - {c}_{n} | \end{align} であるので、この値を最大化するには、

  • なるべく数列  a のうち、大きい値を正の値とする
  • なるべく数列  c のうち、大きい値を正の値とする

ということを同時に達成することが必要です。 従って、数列  c を降順に並び替えることによって、この値を最大化できると考えられます。 すなわち、  {c}_{1} \geq {c}_{2} \geq \cdots \geq {c}_{n} とすることによって最大化できます。

ここで、数列  b から  {a}_{i} 以上の要素を  i 個、  {a}_{i} 未満の要素を  (n - i) 個だけ取り出して数列  c を作成したと考えます。 このとき、  {c}_{1} \geq {c}_{2} \geq \cdots \geq {c}_{i} \geq {a}_{i} であり、  {a}_{i} \lt {c}_{i + 1} \leq \cdots \leq {c}_{n} が成立することから、 \begin{align} D & = | {a}_{1} - {c}_{1} | + | {a}_{2} - {c}_{2} | + \cdots + | {a}_{n} - {c}_{n} | \\ & = ({c}_{1} - {a}_{1}) + \cdots + ({c}_{i} - {a}_{i}) + ({a}_{i + 1} - {c}_{i + 1}) + \cdots + ({a}_{n} - {c}_{n}) \end{align} が成立します。

このことから、  {c}_{1}, \, {c}_{2}, \, \cdots , \, {c}_{i} については正の値を取り、  {c}_{i + 1}, \, {c}_{i + 2}, \, \cdots , \, {c}_{n} については負の値をとります。

従って、  D の値を最大化するためには、数列  b のうち大きい方から  i 個分のものを  {c}_{1}, \, \cdots , \, {c}_{i} とし、小さい方から  (n - i) 個分のものを  {c}_{n}, \, \cdots , \, {c}_{i + 1} とすることによって、達成することができます。

以上から、数列  b を降順に並び替えることを考えます。 すなわち、  {b}_{1} \geq {b}_{2} \geq \cdots \geq {b}_{m} とします。

このとき、数列  c [ {b}_{1}, \, \cdots , \, {b}_{i}, \, {b}_{i + m - n} , \, \cdots , \, {b}_{m} ] とすることによって、  D の値を最大化できます。

これを  0 \leq i \leq n であるすべての  i に対して数列  c を作り、それぞれの場合での  D の値を計算することによって得られる最大値が、この問題の答えとなります。

ただし、毎回数列  c を作成し  D を計算していくと、全体で  O( {n}^{2} ) の計算量がかかってしまい、これは実行時間制限には間に合いません。

そこで、値  \mathrm{sum 0}_{i}, \, \mathrm{sum 1}_{i} をそれぞれ、 \begin{align} & \mathrm{sum 0}_{i} = | {a}_{1} - {b}_{1} | + | {a}_{2} - {b}_{2} | + \cdots + | {a}_{i} - {b}_{i} | \\ & \mathrm{sum 1}_{i} = | {a}_{1} - {b}_{m - n + 1} | + | {a}_{2} - {b}_{m - n + 2} | + \cdots + | {a}_{i} - {b}_{m - n + i}| \end{align} として定めます。 また、  \mathrm{sum 0}_{0} = 0, \, \mathrm{sum 1}_{0} = 0 とします。

この値を整理すると、 \begin{align} & | {a}_{1} - {b}_{1} | + \cdots + | {a}_{i} - {b}_{i} | + | {a}_{i + 1} - {b}_{i + m - n}| + \cdots + | {a}_{n} - {b}_{m} | \\ = & \, \mathrm{sum 0}_{i} + ( | {a}_{1} - {b}_{m - n + 1} | + \cdots + | {a}_{n} - {b}_{m}|) - ( | {a}_{1} - {b}_{m - n + 1} | + \cdots + | {a}_{i} - {b}_{m - n + i} | ) \\ = & \, \mathrm{sum 0}_{i} + \mathrm{sum 1}_{n} - \mathrm{sum 1}_{i} \end{align} によって求めることができます。

従って、  \mathrm{sum 0}_{i}, \, \mathrm{sum 1}_{i} i = 1, \, 2, \, \cdots , \, n まで予め求めておくことにより、各  i に対して  D の値を  O(1) の計算量で求めることができます。 すなわち、全体で  O(n) の計算量で  D の値の最大値を求めることができます。

以上から、数列  a, \, b のソートも含めると、全体で  O(n \log n + m \log m) の計算量で  D の最大値を求めることができ、これは実行時間制限に十分間に合うと考えられます。

ソースコード

各テストケースに対して、  D の値の最大値を返す関数 solve() を以下のように実装しました。

int n, m;
ll a[int(2e5 + 5)], b[int(2e5 + 5)];
ll sum[int(2e5 + 5)][2];

ll solve() {
  // 値の入力
  cin >> n >> m;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  sort(a + 1, a + n + 1); // 数列 a の昇順ソート
  for (int i = 1; i <= m; i++) {
    cin >> b[i];
  }
  sort(b + 1, b + m + 1, greater<ll>()); // 数列 b の降順ソート

  // i = 1, ..., n の順に、上記の sum0, sum1 の値を計算していく
  for (int i = 1; i <= n; i++) {
    sum[i][0] = sum[i - 1][0] + abs(a[i] - b[i]); // sum0 の値を計算する
    sum[i][1] = sum[i - 1][1] + abs(a[i] - b[i + m - n]); // sum1 の値を計算する
  }

  ll ans = 0; // 最終的に答えとなる値 ans

  // i = 0, ..., n に対して D の値を計算していく
  for (int i = 0; i <= n; i++) {
    ans = max(ans, sum[i][0] + (sum[n][1] - sum[i][1])); // max を取り、 ans を更新していく
  }
  return ans;
}

感想

記事を書くにあたって、証明を行いづらい問題でした。