Codeforces Round #545 (Div. 1) B. Camp Schedule

Problem Link: https://codeforces.com/contest/1137/problem/B

There are a lot of ways to solve this problem. For example, we can use Z-algorithm, rolling hash and suffix array (suffix array might be too slow for this constraints). This problem also can be solved by the prefix function. The prefix function appears in KMP algorithm [1].

The prefix function P[k] is defined as the maximum length of a string which is both a prefix and a suffix of S[0 .. k). In this problem, we want to know P[M]. The prefix function can be computed in O(N) time, but I do not explain it here. If you want to know more details, please refer to the book [1].

[1] Thomas H. Cormen et al., Introduction to algorithms, Chapter 32, third edition, The MIT Press, 2009.

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>

using namespace std;

vector<int> kmp(string s) {
  int n = s.size();
  vector<int> p(n + 1);
  p[0] = -1;
  int j = -1;
  for (int i = 0; i < n; i++) {
    while (j != -1 && s[i] != s[j]) j = p[j];
    p[i + 1] = ++j;
  }
  return p;
}

int main() {
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
  string S, T;
  cin >> S >> T;
  int N = S.size();
  int M = T.size();
  vector<int> C(128);
  for (int i = 0; i < N; i++) {
    C[S[i]]++;
  }
  vector<int> P = kmp(T);
  string ans;
  int k = 0;
  while (C[T[k]] > 0) {
    C[T[k]]--;
    ans += T[k];
    k++;
    if (k == M) k = P[k];
  }
  ans += string(C['0'], '0');
  ans += string(C['1'], '1');
  cout << ans << endl;
}