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

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

E - Romantic Glasses / Codeforces Round 918 (Div. 4)

問題

 n 個のグラスが1列に並んでおり、  i 番目のグラスにはジュースが  {a}_{i} 単位だけ入っている。

このグラスの連続する一部分を選択し、奇数番目のグラスに入ったジュースの量の合計と偶数番目のグラスに入ったジュースの量の合計を等しくできるか判定せよ。

すなわち、  1 \leq l \leq r \leq n である整数  l, r を選んだときに、

  •  l, r の偶奇が等しいならば、  {a}_{l} + {a}_{l + 2} + {a}_{l + 4} + \cdots + {a}_{r} = {a}_{l + 1} + {a}_{l + 3} + \cdots + {a}_{r - 1}
  •  l, r の偶奇が異なるならば、   {a}_{l} + {a}_{l + 2} + {a}_{l + 4} + \cdots + {a}_{r - 1} = {a}_{l + 1} + {a}_{l + 3} + \cdots + {a}_{r}

を満たすような  l, r が存在するかを判定せよ。

入力

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

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

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

条件

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

出力

 t 行にわたり、各テストケースについて、等しくできるならば YES を、できないならば NO を出力すること。

解法

まず、整数  l, r を愚直に  l = 1, r = 1 から順番に調べていき、条件を満たす  l, r が存在するのかを調べていく方法が考えられます。 この方法では  O({n}^{2}) の計算量がかかります。

しかし、問題の条件より、  n \leq 2 \cdot {10}^{5} であるので、これでは実行時間制限に間に合いません。

そこで、与えられた条件式を変形して、使いやすい形にできないかを考えます。

与えられた数列  a に対して、条件を満たす整数  l, r が存在すると仮定します。 すなわち、 \begin{align} {a}_{l} + {a}_{l + 2} + \cdots = {a}_{l + 1} + {a}_{l + 3} + \cdots \end{align} がこのとき成立します。 この式を変形すると、 \begin{align} {a}_{l} - {a}_{l + 1} + {a}_{l + 2} - {a}_{l + 3} + \cdots + {(-1)}^{r - l} \cdot {a}_{r} = 0 \end{align} が成立すると言えます。 さらに、この式の符号について整理すると、 \begin{align} {a}_{l} + {(-1)} \cdot {a}_{l + 1} + {(-1)}^{2} \cdot {a}_{l + 2} + {(-1)}^{3} \cdot {a}_{l + 3} + \cdots + {(-1)}^{r - l} \cdot {a}_{r} = 0 \end{align} であり、両辺に  {(-1)}^{l} をかけると、 \begin{align} {(-1)}^{l} \cdot {a}_{l} + {(-1)}^{l + 1} \cdot {a}_{l + 1} + {(-1)}^{l + 2} \cdot {a}_{l + 2} + {(-1)}^{l + 3} \cdot {a}_{l + 3} + \cdots + {(-1)}^{r} \cdot {a}_{r} = 0 \end{align} が成立します。 ここで、左辺の値に着目すると、 \begin{align} \, & {(-1)}^{l} \cdot {a}_{l} + {(-1)}^{l + 1} \cdot {a}_{l + 1} + {(-1)}^{l + 2} \cdot {a}_{l + 2} + {(-1)}^{l + 3} \cdot {a}_{l + 3} + \cdots + {(-1)}^{r} \cdot {a}_{r} \\ = \, & \{ {(-1)} \cdot {a}_{1} + \cdots + {(-1)}^{r} \cdot {a}_{r} \} - \{ {(-1)} \cdot {a}_{1} + \cdots + {(-1)}^{l - 1} \cdot {a}_{l - 1} \} \end{align} と変形することができます。

このことから、値  \mathrm{sum}_{i} を、 \begin{align} \mathrm{sum}_{i} = {(-1)} \cdot {a}_{1} + \cdots + {(-1)}^{i} \cdot {a}_{i} \end{align} として定義すると、問題の条件を満たす  l, r が存在するとき、 \begin{align} \mathrm{sum}_{r} - \mathrm{sum}_{l - 1} = 0 \end{align} より、 \begin{align} \mathrm{sum}_{l - 1} = \mathrm{sum}_{r} \end{align} が成立すると言えます。

ここで、  l = 1 にて問題の条件を満たす  r が存在するときについては、先程の変形から \begin{gather} {(-1)}^{1} \cdot {a}_{1} + {(-1)}^{2} \cdot {a}_{2} + \cdots + {(-1)}^{r} \cdot {a}_{r} = 0 \\ \therefore \; \mathrm{sum}_{r} = 0 \end{gather} ということが言えます。 従って、便宜的に  \mathrm{sum}_{0} = 0 として扱うと、この場合でも  \mathrm{sum}_{l - 1} = \mathrm{sum}_{r} が成立します。

以上から、与えられた数列  a に対して、  \mathrm{sum}_{0}, \mathrm{sum}_{1}, \cdots , \mathrm{sum}_{n} の値を求め、これらの中から等しい値を取るものが1組でもある場合は、条件を満たす  l, r が存在すると判定することができます。

 \mathrm{sum} の値については、 \begin{align} \mathrm{sum}_{i + 1} = \mathrm{sum}_{i} + {(-1)}^{i + 1} \cdot {a}_{i + 1} \end{align} によって求めることができるため、全体で  O(n) の計算量で求めることができます。

また、等しい値があるかの判定については、例えば得られた  \mathrm{sum} を昇順にソートすることで、要素の前後で等しい組があるかを調べ上げて判定することができます。 これはソートに  O(n \log n) の計算量がかかりますが、実行時間制限にも十分間に合うと考えられます。

この問題は  l, r の値を出力する問題ではないので、以上の方法により条件内で判定することができます。

ソースコード

整数  l, r が存在するかを判定し、存在するならば true を、しないならば false を返す関数 solve() を以下のように実装しました。

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

bool solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
    if (i % 2 != 0) {
      sum[i] = sum[i - 1] - a[i]; // i が奇数ならば a[i] を引く
    } else {
      sum[i] = sum[i - 1] + a[i]; // i が偶数ならば a[i] を足す
    }
  }
  sort(sum, sum + n + 1); // sum[0] から sum[n] までを昇順にソートする

  for (int i = 0; i <= n - 1; i++) {
    if (sum[i] == sum[i + 1]) { // sum の値が前後で等しければ、 true を返す
      return true;
    }
  }

  return false; // 前後で等しいものが存在しなければ、 false を返す
}