ABC 161 | D - Lunlun Number

3 min read読了の目安(約3300字

問題

https://atcoder.jp/contests/abc161/tasks/abc161_d

考えたこと

DPを使ってできる。dp[i][j]を最上位がiで桁数がj-1のときのルンルン数の数とする。
ここではi = 0のときも計算する必要あがる。

まず1桁しかないときを考えると以下のようになる。

1 \leq kなのと最上位が0の桁は数えることが赤の部分は数えず、k \leq 9のときは答えはkである。

桁数が2桁のときは以下のようになる。
注意として、最上位が0のときもひとつ下の桁のルンルン数の数を足す。

最上位iが3で桁数jが2のとき、ルンルン数は桁数が1のときのルンルン数のdp[2][0] + dp[3][0] + dp[4][0]で求められる。これは2桁目が3で1桁目が2, 3, 4に対応して実際のルンルン数は32, 33, 34が対応する。

このようにして最上位が0の場合を除いて、桁ごとにそれまでのルンルン数の合計がkに達するまで計算していく。

たとえばk = 47のとき、100の桁が1までのルンルン数の合計が43で、100の桁が2のとき、ルンルン数の合計が52なので100の桁が2のところにk = 47番目のルンルン数がある。そのルンルン数がどの数かを求めるために以下のようにたどっていけばいい。

47 - 43 = 4なのでk = 47のときのルンルン数は100の桁が2のときの4番目の数となる。10の桁をみると、1のときに3つ、2のときに3つなので、10の桁は2だということがわかる。同様に1の桁は1ということがわかるので答えは221である。

コード

#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

// 現在のi, jの値からターゲットとなる値を探す
ll find(vector<vector<ll>> &r, int ci, int cj, int ttl, int target) {
  ll ans = 0;
  target -= ttl;
  for (int j = cj; j >= 0; j--) {
    // 現在の値を足す
    ans += ci * pow(10, j);
    if (j == 0) break;
    ll v = 0;
    for (int i = max(0, ci - 1); i <= min(9, ci + 1); i++) {
      if (v < target && target <= v + r[i][j - 1]) {  // 正しいのをみつけた
        target -= v;
        ci = i;
        break;
      }
      v += r[i][j - 1];
    }
  }
  return ans;
}

int main() {
  ll target;
  cin >> target;
  // leading zeroもあったと仮定したときのルンルン数の数を記録する
  vector<vector<ll>> r(10, vector<ll>(10));  // 100000の答えが3234566667なので10まででよい

  ll ttl = 0;                       //現在のルンルン数
  for (int j = 0; j < 10; j++) {    // 1, 10, 100, ...
    for (int i = 0; i < 10; i++) {  // 0~9
      if (j == 0) {
        r[i][j] = 1;
      } else {
        // 1つ桁が低い周辺を足す
        for (int k = max(0, i - 1); k <= min(9, i + 1); k++) {
          r[i][j] += r[k][j - 1];
        }
      }
      if (i == 0) continue;
      if (target <= ttl + r[i][j]) {
        // 見つかったので実際の値を探す
        cout << find(r, i, j, ttl, target) << endl;
        return 0;
      }
      ttl += r[i][j];
    }
  }
}

別解

ACした後に回答をみたところ最上位の桁をqueueに入れながら計算するやりかたがあった。
以下のように1桁のときの桁をQueueに入れていき、順に取り出し接続可能な数字をつけてQueueにいれていくのである。

こちらのほうが簡単に実装でき、実装間違えが少ない。

#include <bits/stdc++.h>

#include <atcoder/all>

using namespace std;
using namespace atcoder;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
const int MOD = 1e9 + 7;

int main() {
  ll k;
  cin >> k;
  if (k <= 9) {
    cout << k << endl;
    return 0;
  }
  queue<int> q;
  for (int i = 1; i <= 9; i++) {
    q.push(i);
  }
  ll cur = 9;
  while (true) {
    ll v = q.front();
    q.pop();
    ll d = v % 10;
    v *= 10;
    for (int i = max(0ll, d - 1); i <= min(9ll, d + 1); i++) {
      cur++;
      if (cur == k) {
        cout << v + i << endl;
        return 0;
      }
      q.push(v + i);
    }
  }
}

参考

https://atcoder.jp/contests/abc161/tasks/abc161_d/editorial