SRM672 div.1 250 Procrastination

Procrastination.py

# -*- coding: utf-8 -*-
import math,string,itertools,fractions,heapq,collections,re,array,bisect

class Procrastination:
    def findFinalAssignee(self, n):
        def miller_rabin(n):
            """ primality Test
                if n < 3,825,123,056,546,413,051, it is enough to test
                a = 2, 3, 5, 7, 11, 13, 17, 19, and 23.
                Complexity: O(log^3 n)
            """
            assert(n >= 1)
            if n == 2:
                return True
            if n <= 1 or not n & 1:
                return False

            primes = [2, 3, 5, 7, 11, 13, 17, 19, 23]

            d = n - 1
            s = 0
            while not d & 1:
                d >>= 1
                s += 1

            for prime in primes:
                if prime >= n:
                    continue
                x = pow(prime, d, n)
                if x == 1:
                    break
                for r in xrange(s):
                    if x == n - 1:
                        break
                    if r + 1 == s:
                        return False
                    x = x * x % n
            return True

        def divisorsList(n):
            assert(n >= 1)
            divsList = []
            for i in xrange(1, int(n**0.5)+1):
                if n % i == 0:
                    divsList.append(i)
                    if n // i != i:
                        divsList.append(n//i)
            return divsList

        first, last = n, n
        # find prime
        while first >= 2 and not miller_rabin(first - 1):
            first = first - 1
        while not miller_rabin(last):
            last = last + 1

        # find divisors
        workers = range(first, last + 2)
        divlist = collections.defaultdict(list)
        for worker in workers[:-1]:
            for d in divisorsList(worker):
                if d == 1 or d == worker:
                    continue
                else:
                    divlist[d].append(worker)

        # simulation
        for k in sorted(divlist.iterkeys()):
            for v in divlist[k]:
                v = v - first
                workers[v], workers[v+1] = workers[v+1], workers[v]

        return workers[n-first]