ABC192 D - Base n を Python で解いてみる

問題概要

数字列 $X$ を $n$ 進法表記と見なしたとき, $M$ 以下になるような値が何種類あるか求めよ.

ただし, $n$ は $X$ に含まれる数字より大きいものとする.

考えたこと

入力例 1

22
10

22 を 3 進数と見なすと, $2\times 3 + 2 = 8 \leqq 10$
22 を 4 進数と見なすと, $2\times 4 + 2 = 10 \leqq 10$
22 を 5 進数と見なすと, $2\times 5 + 2 = 12 > 10$
22 を 6 進数と見なすと, $2\times 6 + 2 = 14 > 10$

により, 最初の2つだけが条件を満たします.

このことから, $n$ が小さいときには条件を満たし, $n$ を大きくするにつれて $M$ を超えてしまう閾値を見つければよさそうです.

以下, $X$ の長さが $K$ であるとし, $k$ 桁目を $a_k$ ( $a_k\leqq d$ )と表します. このとき, $X$ を $n(>d)$ 進法表記と見なすと,

[tex: \displaystyle X_n = \sum_{k=0}^{K-1} n^k a_k ]

と表すことができ, これは $X$ が二桁以上の場合には, $n$ について単調増大であり, したがっていつかは $M$ よりも大きくなります.

一方で, $X$ が一桁の場合には, 何進法表記であっても

$\displaystyle X_n = a_0$

です. いくら $n$ を大きくしても $M$ を超えないことがあり得ますが, 問題文で問われているのは $M$ 以下になるような値が何種類あるかであるため, $X$ そのものが $M$ 以下であれば $1$ が, そうでなければ $0$ が解答になります

解いていく

まずは, 入力を受け取ります. 数字列 $X$ は整数として入力されますが, $n$ 進数展開を行うことを見越してリストとして受け取りましょう.

X = [int(k) for k in list(input())]
d = max(X)
M = int(input())

また, 数字列 $X$ を $n$ 進数展開したときの値を p_adic(X, n) により計算します.

def p_adic(X, p):
  curr = 0
  for k in X:
    curr += k
    curr *= p
  curr //= p
  return curr

以下のようなコードだと, TLEしてしまいます

for ループが $O(M)$ の例 1

ans = 0
for j in range(d+1, M+1):
  if p_adic(X,j) <= M:
    ans += 1

print(ans)

for ループが $O(M)$ の例 2

ans = 0
for j in range(d+1, M+1):
  if p_adic(X,j) <= M:
    ans = j - d
  else:
    break

print(ans)

今回の制約では $M \leqq 10^{18}$ のため, より高速なアルゴリズムを構築する必要があります.

二分探索

今回の問題では $n$ が小さいときには条件を満たし, $n$ を大きくするにつれて, ある閾値を超えると条件を満たさなくなるという単調性があります.

したがって, その閾値を探す問題に帰着することができます. その際には,

は条件を満たすか？
- 満たさないならば閾値は $0$ から $M/2$ の間にある
- 満たすならば閾値は $M/2$ から $M$ の間にある

という性質を使って, 条件の範囲を1ステップごとに半分にしていくことができます. 効率よく探すことで, $O(\log M)$ の計算量で以上の閾値を求めていきましょう.

一般の二分探索については qiita.com

特に Python の二分探索をうまく書く方法については www.forcia.com

がそれぞれ読みやすく, 参考になります. ポイントとしては,

左端が True, 右端が False となるように初期値 ok, ngを設定する
中央値 (ok + ng)//2 での条件を判定する
- 左端が True, 右端が False となる状態を維持しながら ok, ngを変更する
ok と ng の差が1になったとき, ok の値が「条件を満たす $j$ の上限」である

ことになります.

解いていく

以上を踏まえて, - $X$ が一桁の場合がコーナーケースになっていることに注意する - 「 $X$ を $j$ 進数展開した結果が $M$ 以下である」をis_ok(j)と表記し, 二分探索で最大の $j$ を求める - 答えとして $j - d$ を出力する

を行うことにより, ACすることができます

X = [int(k) for k in list(input())]
d = max(X)
M = int(input())

def p_adic(X, p):
  curr = 0
  for k in X:
    curr += k
    curr *= p
  curr //= p
  return curr

def is_ok(j):
  curr = p_adic(X, j)
  if curr <= M:
    return True
  else:
    return False

if len(X) == 1:
  if X[0] <= M:
    print(1)
  else:
    print(0)
  import sys
  sys.exit()

ok = d
ng = M + 1
while ng - ok > 1:
  mid = (ok + ng) // 2 # 平均(小数切り捨て)
  if is_ok(mid):
    ok = mid
  else:
    ng = mid

print(ok-d)

atcoder.jp