MojoでProject Euler 98

https://projecteuler.net/problem=98

まずアナグラムを見つけます。そして、それを符号化します。例えば、CAREなら、同じ文字を同じ数字で表すとして、9876とします。0を頭につけたくないので、9から始めて、Aは8、Rは7などとします。そのアナグラムのRACEは7896になり、9876と7896のペアを保存します。
次に、平方数を同じように符号化します。1296を符号化して9876、これを元に9216を符号化すると、7896になって、9876と7896がペアになります。このときの7896側の数の最大を辞書に保存しておけばよいです。

import sys


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

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

fn digits_to_number(ds: List[Int]) -> Int:
    var n = 0
    for d in ds:
        n = n * 10 + d[]
    return n

fn join(v: List[String], delim: StringLiteral) -> String:
    var s = String()
    if len(v) == 1:
        return v[0]
    else:
        s = v[0];
    for i in range(1, len(v)):
        s += delim + v[i]
    return s


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

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 reverse_list[T: CollectionElement](a: List[T]) -> List[T]:
    var rev = List[T](capacity=len(a))
    for i in range(len(a)-1, -1, -1):
        rev.append(a[i])
    return rev

fn max_list[T: ComparableCollectionElement](a: List[T]) -> T:
    var max_elem = a[0]
    for i in range(1, len(a)):
        if a[i] > max_elem:
            max_elem = a[i]
    return max_elem

trait Printable(CollectionElement, Stringable):
    pass

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("[]")


#################### Pair ####################

trait HashableElement(ComparableCollectionElement, Hashable):
    pass

@value
struct Pair[S: HashableElement, T: HashableElement](KeyElement):
    var x: S
    var y: T
    
    fn __hash__(self) -> Int:
        return hash(self.x) + hash(self.y)
    
    fn __eq__(self, other: Self) -> Bool:
        return self.x == other.x and self.y == other.y
    
    fn __ne__(self, other: Self) -> Bool:
        return self.x != other.x or self.y != other.y


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

fn read_words(path: String) raises -> List[String]:
    with open(path, "r") as f:
        var data = f.read()
        var v = data.split(",")
        for s in v:
            s[] = s[].replace('"', "")
        return v

fn word_pattern(word: String) -> Int:
    # CARE -> 9876
    var ps = List[Int]()
    var dic = Dict[String, Int]()
    for i in range(len(word)-1, -1, -1):
        var p = dic.get(word[i], 9 - len(dic))
        ps.append(p)
        dic[word[i]] = p
    return digits_to_number(ps)

# 'RACE', 'CARE', 9876 -> 7896
fn word_pattern_by_code(word: String, word_base: String, code_base: Int) -> Int:
    var ds_base = digits(code_base)
    var dic = Dict[String, Int]()
    var L = len(ds_base)
    for i in range(L):
        dic[word_base[i]] = ds_base[L-i-1]
    
    var n = 0
    for i in range(L):
        var d = dic.get(word[i], 0)
        n = n * 10 + d
    return n

fn find_anagrams(words: List[String]) -> Set[Pair[Int, Int]]:
    fn encode(word: String) -> String:
        var cs = List[String]()
        for i in range(len(word)):
            cs.append(word[i])
        sort(cs)
        return join(cs, "")
    
    var dic = Dict[String, List[String]]()
    for word in words:
        var code = encode(word[])
        if code in dic:
            try:
                dic[code].append(word[])
            except:
                pass
        else:
            dic[code] = List[String](word[])
    
    var patterns = Set[Pair[Int, Int]]()
    for v in dic.values():
        var L = len(v[])
        if L == 1:
            continue
        for i in range(L):
            for j in range(L):
                if i == j:
                    continue
                var code1 = word_pattern(v[][i])
                var code2 = word_pattern_by_code(v[][j], v[][i], code1)
                patterns.add(Pair[Int, Int](code1, code2))
    return patterns^

fn max_length(words: List[String]) -> Int:
    var max_len = 0
    for word in words:
        if len(word[]) > max_len:
            max_len = len(word[])
    return max_len

fn encode(n: Int) -> Int:
    var ds = digits(n)
    sort(ds)
    return digits_to_number(reverse_list(ds))

fn classify_squares(max_len: Int) -> Dict[Int, List[Int]]:
    # { 961: [169, 196, 961], ... }
    var dic = Dict[Int, List[Int]]()
    var U = 10**max_len
    for n in range(1, 10**((max_len+1)//2)):
        var sq = n * n
        if sq >= U:
            break
        var m = encode(sq)
        var v = dic.get(m)
        if v is not None:
            try:
                dic[m].append(sq)
            except:
                pass
        else:
            dic[m] = List[Int](sq)
    return dic

fn digits_pattern(n: Int) -> Int:
    # 144 -> 988
    # 1296 -> 9876
    var ds = digits(n)
    var ps = List[Int]()
    var dic = Dict[Int, Int]()
    for i in range(len(ds)-1, -1, -1):
        var p = dic.get(ds[i], 9 - len(dic))
        ps.append(p)
        dic[ds[i]] = p
    return digits_to_number(ps)

# 9216, 1296, 9876 -> 7896
fn num_pattern_by_code(num: Int, num_base: Int, code_base: Int) -> Int:
    var ds1 = digits(num)
    var ds2 = digits(num_base)
    var ds_base = digits(code_base)
    var dic = Dict[Int, Int]()
    var L = len(ds_base)
    for i in range(L):
        dic[ds2[i]] = ds_base[i]
    
    var n = 0
    for i in range(L-1, -1, -1):
        var d = dic.get(ds1[i], 0)
        n = n * 10 + d
    return n

fn find_square_anagrams(max_len: Int) -> Dict[Pair[Int, Int], Int]:
    var cls = classify_squares(max_len)
    var dic = Dict[Pair[Int, Int], Int]()
    for item in cls.items():
        var key = item[].key
        var ns = item[].value
        var L = len(ns)
        for i in range(L):
            for j in range(L):
                if i == j:
                    continue
                var code1 = digits_pattern(ns[i])
                var code2 = num_pattern_by_code(ns[j], ns[i], code1)
                var n = dic.get(Pair[Int, Int](code1, code2), 0)
                if ns[j] > n:
                    dic[Pair[Int, Int](code1, code2)] = ns[j]
    
    return dic

fn max_pair_length(anagrams: Set[Pair[Int, Int]]) -> Int:
    var max_len = 0
    for pair in anagrams:
        var ds = digits(pair[].x)
        if len(ds) > max_len:
            max_len = len(ds)
    return max_len

fn f(path: String) raises -> Int:
    var words = read_words(path)
    var anagrams = find_anagrams(words)
    var max_len = max_pair_length(anagrams)
    var num_anagrams = find_square_anagrams(max_len)
    
    var max_num = 0
    for pair in anagrams:
        if pair[] in num_anagrams:
            var num = num_anagrams.get(pair[], 0)
            if num > max_num:
                max_num = num
    return max_num

fn main() raises:
    print(f("0098_words.txt"))