contest.py

#!/bin/python


import math

def smartsub(N, k):
    if k == 0:
        return 1
    """ k > 1 """
    d, m = divmod(N, 10**k)
    zeroNine = smartsub(10**k-1, k-1)
    if d == 0:
        return smartsub(m, k-1) 
    elif d==1:
        return zeroNine + m +  smartsub(m, k-1) 
    else:
        """ d > 1"""
        return zeroNine *d + 10**k + smartsub(m, k-1)

def smart(N):
    k = math.floor(math.log10(N)) 
    return smartsub(N, k)

def count(i):
    if i == 0:
        return 0
    d, m = divmod(i, 10)
    if m == 1:
        return 1 + count(d)
    else:
        return count(d)

def naive(N):
    ans = 0
    for i in range(1, N+1):
        ans += count(i)
    return ans

def test():
    assert smart(10) == 2
    assert smart(99) == 20
    assert smart(100) == 21
    assert smart(1000) == 301
    assert smart(1000*1000) == 600001