MojoでProject Euler 47

https://projecteuler.net/problem=47

4つの異なる素数からなる連続する4つの自然数を見つけるなら、4つの異なる素数からなる自然数を昇順に生成すればよさそうですが、4つ並んでいるとなるとそれなりにそういう自然数の密度が高くなって、あまり速くなりません。5つだとふつうに篩で素因数分解するよりは速かったものの、1分以上かかってしまいました。

import sys

#################### library ####################

fn make_prime_table(N: Int) -> List[Int]:
    var a = initialize_list(N+1, True)
    for p in range(2, N+1):
        if p * p > N:
            break
        elif not a[p]:
            continue
        
        for n in range(p*p, N+1, p):
            a[n] = False
    
    var ps = List[Int]()
    for n in range(2, N+1):
        if a[n]:
            ps.append(n)
    return ps


#################### List ####################

trait Printable(CollectionElement, Stringable):
    pass

fn initialize_list[T: CollectionElement](N: Int, init: T) -> List[T]:
    var a = List[T](capacity=N)
    for n in range(N):
        a.append(init)
    return a

fn range_list(first: Int, last: Int, delta: Int) -> List[Int]:
    var v = List[Int]()
    if (first < last and delta > 0) or (first > last and delta < 0):
        for n in range(first, last, delta):
            v.append(n)
    return v

fn print_list[T: Printable](a: List[T]):
    if a.size > 0:
        var s = "[" + str(a[0])
        for i in range(1, a.size):
            s += ", " + str(a[i])
        s += "]"
        print(s)
    else:
        print("[]")


#################### HeapQueue ####################

struct HeapQueue[T: ComparableCollectionElement]:
    var a: List[T]
    
    fn __init__(inout self, a: List[T]):
        self.a = a
    
    fn is_empty(self) -> Bool:
        return len(self.a) == 0
    
    fn pop(inout self) -> T:
        var top = self.a[0]
        self.a[0] = self.a.pop()
        var N = self.a.size
        var i = 0
        while i < N-1:
            var j1 = i*2+1
            var j2 = i*2+2
            var j: Int = 0
            if j1 >= N:
                break
            elif j2 == N:
                j = j1
            else:
                if self.a[j1] <= self.a[j2]:
                    j = j1
                else:
                    j = j2
            if self.a[j] < self.a[i]:
                self.swap(i, j)
                i = j
            else:
                break
        
        return top
    
    fn push(inout self, e: T):
        self.a.append(e)
        var i = self.a.size - 1
        while i > 0:
            var j = (i-1) // 2
            if self.a[j] <= self.a[i]:
                break
            else:
                self.swap(i, j)
                i = j
    
    fn swap(inout self, i: Int, j: Int):
        var tmp = self.a[i]
        self.a[i] = self.a[j]
        self.a[j] = tmp^


#################### Factors ####################

struct Factors(ComparableCollectionElement):
    var ks: List[Int]
    var es: List[Int]
    var value: Int
    
    fn __init__(inout self, ks: List[Int], es: List[Int], value: Int):
        self.ks = ks
        self.es = es
        self.value = value
    
    fn __copyinit__(inout self, other: Factors):
        self.ks = other.ks
        self.es = other.es
        self.value = other.value
    
    fn __moveinit__(inout self, owned other: Factors):
        self.ks = other.ks^
        self.es = other.es^
        self.value = other.value
    
    fn __len__(self) -> Int:
        return len(self.ks)
    
    fn __lt__(self, other: Self) -> Bool:
        return self.value < other.value
    
    fn __le__(self, other: Self) -> Bool:
        return self.value <= other.value
    
    fn __eq__(self, other: Self) -> Bool:
        return self.value == other.value
    
    fn __ne__(self, other: Self) -> Bool:
        return self.value != other.value
    
    fn __gt__(self, other: Self) -> Bool:
        return self.value > other.value
    
    fn __ge__(self, other: Self) -> Bool:
        return self.value >= other.value
    
    fn __mul__(self, other: Factors) -> Factors:
        var ks = List[Int]()
        var es = List[Int]()
        var L1 = self.ks.size
        var L2 = other.ks.size
        var k = 0
        var l = 0
        while k < L1 and l < L2:
            var k1 = self.ks[k]
            var e1 = self.es[k]
            var k2 = other.ks[l]
            var e2 = other.es[l]
            if k1 == k2:
                ks.append(k1)
                es.append(e1+e2)
                k += 1
                l += 1
            elif k1 < k2:
                ks.append(k1)
                es.append(e1)
                k += 1
            else:
                ks.append(k2)
                es.append(e2)
                l += 1
        
        for k1 in range(k, L1):
            ks.append(self.ks[k1])
            es.append(self.es[k1])
        for l1 in range(l, L2):
            ks.append(other.ks[l1])
            es.append(other.es[l1])
        return Factors(ks, es, self.value * other.value)
    
    fn __imul__(inout self, other: Factors):
        for i in range(len(other.ks)):
            var k = other.ks[i]
            var e = other.es[i]
            for j in range(len(self.ks)):
                if self.ks[j] == k:
                    self.es[j] += e
                    break
            else:
                self.ks.append(k)
                self.es.append(e)
        self.value *= other.value
    
    fn string(self, primes: List[Int]) -> String:
        if len(self) == 0:
            return "1"
        
        var s = str(self.value) + " = " + str(primes[self.ks[0]])
        if self.es[0] > 1:
            s += "^" + str(self.es[0])
        for i in range(1, len(self)):
            s += " * " + str(primes[self.ks[i]])
            if self.es[i] > 1:
                s += "^" + str(self.es[i])
        return s


#################### process ####################

fn is_descending(es: List[Int]) -> Bool:
    for i in range(1, len(es)):
        if es[i-1] < es[i]:
            return False
    return True

fn nexts_descending(fs: Factors, primes: List[Int]) -> List[Factors]:
    var fss = List[Factors]()
    
    # 左端は常に積める
    var es1 = fs.es
    es1[0] += 1
    fss.append(Factors(fs.ks, es1, fs.value * 2))
    
    # 左端と違う指数が前と一つ差なら積める
    for i in range(1, len(fs)):
        if fs.es[i] == fs.es[0] - 1:
            var es2 = fs.es
            es2[i] += 1
            fss.append(Factors(fs.ks, es2, fs.value * primes[i]))
            break
        elif fs.es[i] != fs.es[0]:
            break
    
    return fss

fn nexts_swap(fs: Factors, primes: List[Int]) -> List[Factors]:
    var fss = List[Factors]()
    for k in range(len(fs) - 1, 0, -1):
        if fs.es[k] > fs.es[k-1]:
            if k > 1 and fs.es[k-2] >= fs.es[k]:
                var es1 = fs.es
                es1[k-1] = fs.es[k-2]
                es1[k-2] = fs.es[k-1]
                var de = fs.es[k-2] - fs.es[k-1]
                var value = fs.value // primes[fs.ks[k-2]]**de
                value *= primes[fs.ks[k-1]]**de
                fss.append(Factors(fs.ks, es1, value))
            break
        elif fs.es[k] < fs.es[k-1]:
            var es2 = fs.es
            es2[k] = fs.es[k-1]
            es2[k-1] = fs.es[k]
            var de = fs.es[k-1] - fs.es[k]
            var value = fs.value // primes[fs.ks[k-1]]**de
            value *= primes[fs.ks[k]]**de
            fss.append(Factors(fs.ks, es2, value))
    return fss

fn nexts_shift(fs: Factors, primes: List[Int]) -> List[Factors]:
    # 指数を移動できる
    var fss = List[Factors]()
    
    # 最後は移動できる
    var ks1 = fs.ks
    var back = len(ks1) - 1
    ks1[back] += 1
    var value = fs.value // primes[fs.ks[back]]**fs.es[back]
    value *= primes[fs.ks[back]+1]**fs.es[back]
    fss.append(Factors(ks1, fs.es, value))
    
    # 最後の0が連続していないとき手前を右に一つ移動できる
    for i in range(len(fs.ks)-2, -1, -1):
        if fs.ks[i] == fs.ks[i+1] - 2:
            var ks1 = fs.ks
            ks1[i] += 1
            var value = fs.value // primes[fs.ks[i]]**fs.es[i]
            value *= primes[fs.ks[i]+1]**fs.es[i]
            fss.append(Factors(ks1, fs.es, value))
            break
        elif fs.ks[i] < fs.ks[i+1] - 2:
            break
    
    return fss

fn is_packed(fs: Factors) -> Bool:
    var L = len(fs.ks)
    return fs.ks[L-1] == L-1

fn nexts(fs: Factors, primes: List[Int]) -> List[Factors]:
    var fss = List[Factors]()
    if is_packed(fs):
        if is_descending(fs.es):
            var fss1 = nexts_descending(fs, primes)
            for fs1 in fss1:
                fss.append(fs1[])
        var fss1 = nexts_swap(fs, primes)
        for fs1 in fss1:
            fss.append(fs1[])
    
    var fss2 = nexts_shift(fs, primes)
    for fs2 in fss2:
        fss.append(fs2[])
    return fss

fn f(L: Int, M: Int) -> Int:
    var primes = make_prime_table(1000000)
    var ks_init = range_list(0, L, 1)
    var es_init = initialize_list(L, 1)
    var value = 1
    for k in ks_init:
        value *= primes[k[]]
    var fs_init = Factors(ks_init, es_init, value)
    var pq = HeapQueue[Factors](List[Factors](fs_init))
    var counter = 0
    var prev = 0
    while True:
        var fs = pq.pop()
        if fs.value == prev + 1:
            counter += 1
            if counter == M:
                return fs.value - M + 1
        else:
            counter = 1
        for fs1 in nexts(fs, primes):
            pq.push(fs1[])
        prev = fs.value

fn main() raises:
    var args = sys.argv()
    var L = atol(args[1])
    var M = atol(args[2])
    print(f(L, M))