Code Jam Japan2011 練習問題B. 数の集合 (Small) 解いた

概要

最近の息抜きであるCode Jam Japan 2011の問題B. 数の集合を台風の足止めされている中、夢中になって解いた(Smallのみ)ので、またスクリプトを晒してみる。


ざっくりと問題説明

連続した範囲[A,B]の整数に対して、素因数分解をして指定された値P以上の素数が共通している整数を同じ組にして、最終的に何組になりましたか?という問題。デフォルトの整数N個はそれぞれ自分しか所属していない集合に属している。(デフォルト集合はN個)


Pythonスクリプト

#!/usr/bin/env python
# coding: utf-8
# vim: fileencoding=utf8

import sys
import re
import datetime


def read_input_data(file_name):
        try:
                file_obj = file(file_name, 'r')
                length = int(file_obj.readline().rstrip('\r\n'))
                return True, length, \
                        [map(int, re.compile(u' ').split(line.rstrip('\r\n'))) for line in file_obj]
        except IOError:
                print 'FILE NOT FOUND: ' + str(file_name)
                return False, None, None


def get_prime_factorization(num, prime_list):
        if num in prime_list: return [1, num]
        
        prime_factors = [1,]
        for prime in prime_list[1:]:
                while(True):
                        if num % prime == 0:
                                prime_factors.append(prime)
                                num /= prime
                        else:
                                break
                if num == 1: break
        return prime_factors


def get_prime(max_num):
        prime_list = [1,]
        for candidate in range(2, max_num+1):
                for j in range(2, candidate+1):
                        if candidate == j:
                                prime_list.append(candidate)
                                break
                        if candidate % j == 0:
                                break
                        if (candidate / 2) == j:
                                prime_list.append(candidate)
                                break
        return prime_list


def get_num_duplicated(dup_set_list):
        join_flg = False 
        for i in range(0, len(dup_set_list)):
                for j in range(i+1, len(dup_set_list)):
                        if len(dup_set_list[i] & dup_set_list[j]) > 0:
                                join_flg = True
                                joined_set = dup_set_list[i] | dup_set_list[j]
                                dup_set_list[i] = set() 
                                dup_set_list[j] = joined_set
        if join_flg:
                return get_num_duplicated([value for value in dup_set_list if len(value) > 0])
        return sum(map(len, dup_set_list)) - len(dup_set_list)


if __name__ == '__main__':
        start_time = datetime.datetime.now()
        
        input_file = './B-small-practice.in'
        #input_file = './B-large-practice.in'
        
        read_judge, length, dataset = read_input_data(input_file)
        if not read_judge: sys.exit()
        
        max_num = max([data[1] for data in dataset])
        prime_list = get_prime(max_num)
        
        # for cache: 素因数分解した値を保存するため
        prime_factorization_list = {}
        for i in range(length):
                data_min = dataset[i][0]
                data_max = dataset[i][1]
                data_range = range(data_min,data_max+1)
                p = dataset[i][2]
                p_set = sorted(set(range(p, data_max+1)) & set(prime_list))
                
                join_candidate = {} 
                num_set = len(data_range)
                for num in data_range:
                        # 素因数分解する. 既存のやつはcacheの prime_factorization_list を使用
                        try:
                                prime_factorization_list[num]
                        except KeyError: 
                                prime_factorization_list[num] = set(get_prime_factorization(num, prime_list))
                        
                        # p以上の素因数を持つものを join_candidate に束ねる
                        for p_prime in p_set:
                                if p_prime in prime_factorization_list[num]:
                                        try:
                                                join_candidate[p_prime]
                                        except KeyError:
                                                join_candidate[p_prime] = []
                                        finally:
                                                join_candidate[p_prime].append(num)
                # 2つ以上一致するものを選択
                duplicated_list = [value for value in join_candidate.values() if len(value) > 1]
                answer = num_set - get_num_duplicated(map(set, duplicated_list))
                print "Case #%s: %s" %(i+1, answer)

Largeに適用してみた。

そのままLargeの入力データを指定して実行してみると以下のエラー。

OverflowError: range() result has too many items

などと怒られた。さすがに12桁の数を普通に扱うには無理があった模様。
その問題は、以下のエントリーが解決してくれた。
Python: range and xrange for 13-digit numbers?
スクリプト変更後、実行できた!けど、レスポンスが返ってこない・・・。あれれー。おかしいぞ〜。
printを使ってどこで時間がかかっているか調べたら、初っ端の素数のリストを求めるところで詰んでいた。
いろいろ調べたところ、エラトステネスの篩というアルゴリズムが優秀らしい。その実装は次回にでも。


感想

久々に素数を求めるアルゴリズムを書くことになり、とっさにはパッと書くことができなかった。
また素数をいかに高速に求めるかというのはプログラミングコンテストになるくらいの課題だということに驚いた。
1000万個目の素数を超高速に出力せよ - Team-lablog
しかし、大きい数を扱うのはいろいろと大変だし、癖があることを勉強することができた。Cのコードもセグフォになっちゃうし・・・。JavaBigIntegerクラスだと簡単にできそうだったが、せっかくだし、Pythonでこの問題を乗り越えたいな〜っと思ったので、近いうちにまたチャレンジする予定。
それにしてもCの実行速度の速さにはびっくりした。Cもちゃんと勉強したいな。いや〜プログラミング楽しいっすね。