inputの落とし穴(SPOJ 11. Factorial)

今日のテーマはSPOJの11番Factorial．問題を要約すると，自然数が与えられて，その階乗の値の末尾に何個の0が続くかを求める問題．解法としては階乗の値を素因数分解したときに5がいくつ含まれているかを数えればいい．まずはシンプルに実装してみよう．

#!/usr/bin/env python2.6
def Z(n):
    k = 5
    c = 0
    while k <= n:
        c += n / k
        k *= 5
    return c

num = input()
for i in xrange(num):
    print Z(input())

これでサンプルの出力と一致させることができたのでサブミットしてみる．しかし結果は「time limit exceeded(時間切れ)」．そこで問題をもう一度よく見てみると，入力は1以上1000000000以下の自然数が100000個で，これを6秒以内に計算しなければいけないとのこと．アルゴリズム的に問題がないかどうか確かめるために一度Cで実装してみよう．

#include <stdio.h>

int Z(int x){
    int k = 5;
    int count = 0;
    while(k <= x){
        count += x / k;
        k *= 5;
    }
    return count;
}

int main(void){
    int n, x, i;
    scanf("%d", &n);
    for(i=0; i<n; i++){
        scanf("%d", &x);
        printf("%d\n", Z(x));
    }
}

これでサブミットしたところ，0.25秒程度でアクセプトされた．アルゴリズム的な問題ではなく，Pythonでの実装上の問題ということがわかったが，20倍以上遅いとは…．そこでZ(n)を次のように書きかえてみる．

import math
limit = 1000000000
L = [5**(i+1) for i in range(int(math.log(limit, 5)))]
def Z(n):
    return sum([n/k for k in L])

これで10%程高速化できたが，まだアクセプトされない．これ以上はお手上げだったのでフォーラムを覗いてみた．そこにはなんと，Pythonのinput()はちょっと遅いからsys.stdin.readline()を使えと書いてある．そういうわけで，言われたとおりにしてみよう．

import sys
import math

limit = 1000000000
L = [5**(i+1) for i in range(int(math.log(limit, 5)))]
def Z(n):
    return sum([n/k for k in L])

n = int(sys.stdin.readline())
for i in xrange(n):
    print Z(int(sys.stdin.readline()))

これでようやくアクセプト．

Pythonのinput()は入力された文字列を式として解釈するらしい．つまり三項演算やラムダ式，リスト内包など，式であればなんでも評価して値を返してくれるが，その分処理速度は低下するということか．上記のプログラムでint(sys.stdin.readline())の代わりにint(raw_input())としてもアクセプトされることもわかった．
Python3ではinputが廃止され，raw_inputがinputへと改名されるようだ．Python2でのinputと同じことを3で行うためには，eval(input())とする．