mph.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import typing
import random
import math

from typeguard import typechecked

class MPH(list):

    def __init__(self, a: int, b: int, p: int, *args, **kwargs):
        list.__init__(self, *args, **kwargs)
        self.a = a
        self.b = b
        self.p = p

    def size(self) -> int:
        result = len(self)
        for x in self:
            if isinstance(x, MPH):
                result += x.size()
        return result

    def __repr__(self):
        lines = [
            "MPH {",
            "  a = {}, b = {}, p = {}".format(self.a, self.b, self.p),
            "  [",
            "\n".join(
                [ "\n".join([ "    " + y for y in str(x).split("\n") ])
                    for x in self ]),
            "  ]",
            "}",
        ]
        return "\n".join(lines)

@typechecked
def is_prime_number(x: int) -> bool:
    for i in range(2, int(math.sqrt(x) + 1)):
        if x % i == 0:
            return False
    return True

DEFAULT_UPPER_BOUND = 2 ** 16

@typechecked
def random_prime_number(
        lower_bound: int=2,
        upper_bound: int=DEFAULT_UPPER_BOUND) -> int:
    while True:
        result = random.randrange(lower_bound, upper_bound)
        if is_prime_number(result):
            return result

@typechecked
def default_hash(a: int, b: int, p: int, x: int) -> int:
    return (a * x + b) % p

@typechecked
def make_square_mph(
        s: typing.Container[typing.Any],
        h: typing.Callable[[int, int, int, typing.Any], int]=default_hash,
        upper_bound: int=DEFAULT_UPPER_BOUND
        ) -> MPH:
    n = len(s) * len(s)
    while True:
        a = random.randrange(0, upper_bound)
        b = random.randrange(0, upper_bound)
        p = random_prime_number(n + 1, upper_bound)
        result = MPH(a, b, p, [ None ] * n)
        for x in s:
            retry_needed = False
            i = h(a, b, p, x) % n
            if result[i] is None:
                result[i] = x
            else:
                retry_needed = True
                break
        if retry_needed:
            continue
        return result

@typechecked
def make_random_mph(
        iterable: typing.Iterable[typing.Any],
        h: typing.Callable[[int, int, int, typing.Any], int]=default_hash,
        upper_bound: int=DEFAULT_UPPER_BOUND
        ) -> MPH:
    s = set(iterable)
    n = len(s)
    while True:
        a = random.randrange(0, upper_bound)
        b = random.randrange(0, upper_bound)
        p = random_prime_number(n + 1, upper_bound)
        result = MPH(a, b, p, [ None ] * n)
        for x in s:
            i = h(a, b, p, x) % n
            if result[i] is None:
                result[i] = x
            elif isinstance(result[i], list):
                result[i].append(x)
            else: # collision
                result[i] = [ result[i], x ]
        retry_needed = False
        for i, x in enumerate(result):
            if isinstance(x, list):
                try:
                    result[i] = make_square_mph(x, h, upper_bound)
                except ValueError: # empty range for randrange()
                    retry_needed = True
                    break
        if retry_needed:
            continue
        return result