ABC114 C問題 755を3進数を使って解く

ABC114 C問題 755 を公式解説とは異なる実装方法の解説をします。
公式解説の深さ優先探索を使った書き方の方が簡潔で美しいです。
私の書き方は愚直な感じでスマートではありません。

atcoder.jp

問題

整数 N が与えられます。
1 以上 N 以下の整数のうち、七五三数 は何個あるでしょうか?
ここで、七五三数とは以下の条件を満たす正の整数です。
十進法で表記したとき、数字 7, 5, 3 がそれぞれ 1 回以上現れ、これら以外の数字は現れない。

難しいポイント

Nが最大で 10^9なので、1からNまで何も考えずに七五三数かどうか判定して、数を数えるのはPythonでは制限時間には収まらないです。
あと方針が決まっても実装方法もいくつか思いつくので迷いが生じる気がします。

解き方の方針

数字7, 5, 3が一回ずつ現れる数字はそれほど多くなさそうというところに着目します。
9桁それぞれで7か5か3を選ぶ場合の数を考えても全部で 3^9 = 19683通りです。
ビット全探索の発想を応用して、3進数で全探索して、N以下という条件と7, 5, 3が1回ずつ出現する条件の両方を満たす数を数えます。

実装方法

3進数で、2を7、1を5、0を3に見立てて、全部列挙して七五三数なのか、N以下なのか確認して、数を数えていきます。

3進数を列挙するときにproductメソッドを使いました。
以下の場合は、0, 1, 2から4桁の全ての場合の数を列挙してくれます。

from itertools import product
product(range(3), repeat=4)
# (0, 0, 0, 0), (0, 0, 0, 1) ... (2, 2, 2, 2)

3進数を10進数への変換も必要です。
例えば、210 だったら753に変換します。
3進数の数字を配列としてみた時に、インデックス0は2です。これは7に変換して 10^2をかけます。
配列の長さ - インデックス数 - 1 が10進数で何乗するかの値になります。
3桁の場合を列挙すると以下のようになります。
3 - 0 - 1 = 2 ->  10^2 = 100
3 - 1 - 1 = 1 ->  10^1 = 10
3 - 2 - 1 = 0 ->  10^0 = 1

コードは以下のようになりました。

from math import log10
from itertools import product
N = int(input())

# Nの桁数を求める
max_keta = int(log10(N)) + 1

ans = 0

# 1桁からNの桁数までループする
for keta in range(1, max_keta + 1):
    search_list = product(range(3), repeat=keta)
    for sl in search_list:
        # 3進数で 2 -> 7, 1 -> 5, 0 -> 3 とする
        # 7か5か3が1回以上含まれていなかったら次のパターンへ行く
        if not (2 in sl and 1 in sl and 0 in sl):
            continue
        # 210 だったら 753 に変換する
        # len = 3 のとき index 0 だったら 10^2 をかけて足す
        # len - index - 1 で 3 - 0 - 1 = 2, 3 - 1 - 1 = 1, 3 - 2 - 1 = 0
        number = 0
        TRANS = [3, 5, 7]
        for j, n in enumerate(sl):
            number += 10**(len(sl) - j - 1) * TRANS[n]
        if number <= N:
            ans += 1

print(ans)


公式解説を読んで、深さ優先探索再帰的に処理する方法に慣れないといけないなと思いました。