🕌

ABC220 C - Long Sequence C++解答例

2021/09/27に公開

AtCoder Beginner Contest 220のC問題をC++で解きます。

問題

問題文をAtCoderのページより引用します。

問題文

長さNの正整数のみからなる数列A = (A_1, \ldots, A_N)があります。
A10^{100}回連結した数列を数列Bとします。

Bの項を前から順に足したとき、和が初めてXを超えるのは何項目まで足したときですか?
すなわち、以下の式を満たす最小の整数kを求めてください。

\sum_{i = 1}^{k} B_{i} > X

制約

  • 1 \leq N \leq 10^5
  • 1 \leq A_{i} \leq 10^9
  • 1 \leq X \leq 10^{18}
  • 入力は全て整数

解答例

考え方

愚直に考えるとTLE

数列Bの定義からしてBを直接構築することは不可能であることがわかると思います。まず間違いなくメモリが足りません。

Bを作らずに解こうと思っても、Xの最大値が10^{18}なので、「Xを超えるまでA_{i}の値をどんどん足していく。A_{N}まで到達して尚Xを超えなかったらA_{1}に戻ってまた足していく」という操作を繰り返したとしても、Xは非常に大きいのでループ回数が大きくなることが予想されます。

Bの定義からヒントを得る

数列Bの定義をよく見ると、「数列Aをそのままの順序で複数個結合している」ことが分かります。
これを踏まえると、「A_{1}からA_{N}まで順に足していって尚Xを超えなかったとき、最初に連結されたAを除いた数列BX - S(Sは数列Aの各要素の総和)について探索する」ことになることが分かります。

これを拡張して考えてみます。求める項kに対して、数列B'を次のように定義し直してみます。

B' = (B'_{1}, B'_{2}, \ldots, B'_{k}):数列A全体をq個(0 \leq q)連結したのち、Aの先頭からs個(0 \leq s < N)の項を末尾に連結した数列。

定義より、この数列B'の項の数はq \times N + sとなります。
ということは、このqsが求まれば問題を解くことができます。

では実際にqsを求めてみます。まずはqを求めます。

数列Aq個連結した数列を数列Qと呼ぶことにします。このとき、数列B'の定義より、Qの各要素の総和はX以下になります。
また、sの取りうる範囲を考えると、qの値はできる限り最大の値を取らなければならないことがわかります。

従って、qの値は、XSで割ったときの商であることがわかります。

次にsを求めます。qが求まったことにより、sを求める問題は、XSで割ったときの余りをrとすると、「数列Aの各要素を先頭から順に足していったとき、和が初めてrを超えるのは先頭から何項目か?」というものになります。
この問題を解くために探索しなければならない範囲は数列Aだけです。つまり\mathcal{O}(N)で解くことができます。

以上でqsが求まったので、この問題の解はq \times N + sとなります。

プログラムを実装してみる

以上の考え方を実装したのが以下のコードです。

c.cpp
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <numeric>
 
int main() {
  int64_t N;
  std::cin >> N;
  std::vector<int64_t> A(N);
  for (int64_t i = 0; i < N; i++) {
    std::cin >> A.at(i);
  }
  int64_t X;
  std::cin >> X;

  // 数列Aの総和
  int64_t sum = std::accumulate(A.begin(), A.end(), static_cast<int64_t>(0));
  // XをAの総和で割った商
  int64_t quotient = X / sum;
  // XをAの総和で割った余り
  int64_t residue = X % sum;

  // 数列Aの各要素を先頭から足していったときの和
  int64_t total = 0;
  // 和が初めてrを超えたときの項
  int64_t index = 0;
  for (int64_t i = 0; i < N; i++) {
    total += A.at(i);
    index = i + 1;
    if (total > residue) {
      break;
    }
  }
 
  std::cout << index + quotient * N << std::endl;
 
  return 0;
}

実際に提出したコードはこちら

Discussion