SRM546 Div1 Medium

500

m 桁目を n より大きくして残りは０で埋める。下の桁で digit1, digit2 の帳尻を合わせる。



もう少しスッキリした実装ができても良かった気がする。

vector<lli> l2v(lli n)
{
  vector<lli> v;
  while (n) {
    v.push_back(n % 10LL);
    n /= 10LL;
  }
  reverse(v.begin(), v.end());
  return v;
}

lli v2l(vector<lli> v)
{
  lli n = 0;
  for (int i = 0; i < (int)v.size(); ++i) {
    n = n * 10LL + v[i];
  }
  return n;
}

map<lli, int> count_digit(vector<lli> v)
{
  map<lli, int> cnt;
  bool flg = false;
  for (int i = 0; i < (int)v.size(); ++i) {
    if (v[i]) flg = true;
    if (flg) ++cnt[v[i]];
  }
  return cnt;
}

lli f(lli N, int d1, int c1, int d2, int c2)
{
  vector<lli> v = l2v(N);
  map<lli, int> cnt = count_digit(v);
  if (c1 <= cnt[d1] && c2 <= cnt[d2]) return N;
  return 1LL << 60;
}

lli g(lli N, int d1, int c1, int d2, int c2)
{
  lli ret = 1LL << 60;

  vector<lli> v = l2v(N);
  while (v.size() < 16) v.insert(v.begin(), 0);

  const int pref = v.size();
  for (int i = 0; i < pref; ++i) {
    for (int j = v[i]; j <= 9; ++j) {
      vector<lli> u(v.size(), 0);
      copy(v.begin(), v.begin() + i, u.begin());
      u[i] = j;

      map<lli, int> cnt = count_digit(u);
      reverse(u.begin(), u.end());

      int k = 0, l = 0;
      for (; k < c2 - cnt[d2]; ++k) {
        u[k] = d2;
      }
      for (; l < c1 - cnt[d1]; ++l) {
        u[k + l] = d1;
      }

      reverse(u.begin(), u.end());

      cnt = count_digit(u);
      if (c1 <= cnt[d1] && c2 <= cnt[d2]) {
        lli n = v2l(u);
        if (N <= n) ret = min(ret, v2l(u));
      }
    }
  }

  return ret;
}

class FavouriteDigits {
public:
  long long findNext(long long N, int d1, int c1, int d2, int c2)
  {
    if (d1 < d2) ; else {
      swap(d1, d2);
      swap(c1, c2);
    }

    lli a = f(N, d1, c1, d2, c2);
    lli b = g(N, d1, c1, d2, c2);

    return min(a, b);
  }
};