MojoでProject Euler 34

https://projecteuler.net/problem=34

Problem 30とほぼ同じで重複組み合わせを出します。ただし、0! = 1なので、桁数を変えないといけません。

from algorithm.sort import sort
from math import min
import sys


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

fn print_vector(v: List[Int]):
    for k in range((v.size + 9) // 10):
        var s = str("")
        for i in range(k*10, min(k*10+10, v.size)):
            s += str(v[i]) + " "
        print(s)

fn digits(n: Int) -> List[Int]:
    var m = n
    var ds = List[Int]()
    while m > 0:
        var d = m % 10
        m //= 10
        ds.append(d)
    return ds

fn factorial(n: Int) -> Int:
    return 1 if n == 0 else n * factorial(n-1)


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

fn duplicate_combinations(a: List[Int], n: Int) -> List[List[Int]]:
    var v = List[List[Int]]()
    if n == 1:
        for e in a:
            var w = List[Int](e[])
            v.append(w)
    elif n % 2 == 1:
        var v1 = duplicate_combinations(a, n-1)
        for e in a:
            for w1 in v1:
                if e[] > w1[][0]:
                    continue
                var w = List[Int](e[])
                w.extend(w1[])
                v.append(w)
    else:
        var m = n // 2
        var v1 = duplicate_combinations(a, m)
        for w1 in v1:
            for w2 in v1:
                if w1[][m-1] > w2[][0]:
                    continue
                var w = w1[]
                w.extend(w2[])
                v.append(w)
    return v

fn upper_num_digits() -> Int:
    for n in range(1, 20):
        if n * factorial(9) < 10**(n-1):
            return n - 1
    return 0

fn is_coincident(w: List[Int], v: List[Int]) -> Bool:
    if w.size != v.size:
        return False
    
    for i in range(w.size):
        if w[i] != v[i]:
            return False
    return True

fn valid_sum_factorials(v: List[Int]) -> Int:
    var s = 0
    for e in v:
        s += factorial(e[])
    
    var w = digits(s)
    sort(w)
    
    if is_coincident(w, v):
        return s
    else:
        return 0

fn f() -> Int:
    var N = upper_num_digits()
    var ds = List[Int]()
    for d in range(10):
        ds.append(d)
    var s = 0
    for n in range(2, N+1):
        var v = duplicate_combinations(ds, n)
        for w in v:
            s += valid_sum_factorials(w[])
    return s

fn main() raises:
    print(f())