COLOCON 2018-Final: A - ファイティング・タカハシ

問題

A - ファイティング・タカハシ

解説

'B'で区切られた'A'について、それぞれの'A'の連結成分においてスコアを算出する。

ここでNの値が大きいので、当然SN個連結して解くことはできない。

Sの末尾とSの先頭が連結するかどうかとに着目する必要がある。 またすべてが'A'であるとき、N*|S|個の'A'が存在し、スコアの算出方法が異なるため、別途取り扱う必要がある。

つまり、以下の3つに場合分けできれば良い。

  1. Sのすべての要素が'A'
  2. Sの末尾は'A' かつ Sの先頭が'A'
  3. Sの末尾は'B' または Sの先頭が'B'

'B'は区切り文字としてしか機能しないため、Sを'A'で作られる文字列の配列aに変換し、a_0, a_{n-1}, a_{0...n-2} に分けてスコアの算出を考えた。

auto make_A_vec(string const& S) {
  vector<string> ret;
  string curr;
  for (auto c: S) {
    if (c == 'B') {
      if (curr.size()) {
        ret.push_back(curr);
        curr.clear();
      }
    } else {
      curr += 'A';
    }
  }
  if (curr.size()) {
    ret.push_back(curr);
  }
  return ret;
}

template<typename InputIterator>
auto sum_A(InputIterator first, InputIterator last) {
  int64_t sum = 0;
  while (first != last) {
    sum += first->size() * (first->size() + 1) / 2;
    first++;
  }
  return sum;
}

int main() {
  int N; cin >> N;
  string S; cin >> S;
  {
    auto t = S;
    t.erase(std::unique(t.begin(), t.end()), t.end());
    if (t.size() == 1) {
      if (t[0] == 'B') {
        cout << 0 << "\n";
        return 0;
      }
      int64_t x = S.size() * N;
      cout << x * (x + 1) / 2 << "\n";
      return 0;
    }
  }

  auto v = make_A_vec(S);

  if (S.back() == 'A' && S[0] == 'A') {
    auto X = (int64_t)v[0].size() * (v[0].size() + 1) / 2;
    auto sumY = v.size() > 2 ? sum_A(v.begin() + 1, v.end() - 1) : 0;
    auto Z = (int64_t)v.back().size() * (v.back().size() + 1) / 2;
    auto ZX = (int64_t)(v.back().size() + v[0].size()) * ((v.back().size() + v[0].size()) + 1) / 2;
    cout << (X + sumY) + (N - 1) * (ZX + sumY) + Z << "\n";
  } else {
    int64_t sum = 0;
    for (auto e: v) {
      sum += (int64_t)e.size() * (e.size() + 1) / 2;
    }
    cout << sum * N << "\n";
  }
}

  /*
  // 実装力のなさから単体テストを書く羽目になった
  vector<vector<string>> vs = {
    {"A", "A", "A"},
    {"A", "A", "A", "A"},
    {"A", "A", "AA", "A"},
    {"A", "A", "AAA", "A"},
    {"A", "AA", "AAA", "AAAA", "A"},
  };
  vector<int64_t> res_vs = {
    1,
    2,
    4,
    7,
    3 + 6 + 10
  };
  rep(i, res_vs.size()) {
    assert(sum_A(vs[i].begin() + 1, vs[i].end() - 1) == res_vs[i]);
  }
  return 0;
  */

/*
  assert((make_A_vec("A") == vector<string>{"A"}));
  assert((make_A_vec("AA") == vector<string>{"AA"}));
  assert((make_A_vec("ABA") == vector<string>{"A", "A"}));
  assert((make_A_vec("ABABB") == vector<string>{"A", "A"}));
  assert((make_A_vec("ABABBAA") == vector<string>{"A", "A", "AA"}));
  return 0;
*/