https://projecteuler.net/problem=12
ふつうに三角数を素因数分解しても十分速いです。
しかし、素因数分解の指数だけで約数の個数は決まるので、約数の個数が条件より大きい指数だけ考えれば速いです。
素因数分解を
とすると、約数の個数は、
となります。ゆえに、eを並び替えたり飛ばしたりしても約数の個数は同じです。例えば、と
は約数の個数は同じですが、前者のほうが数としては小さいです。一般に、素数が飛ばされていなくて、指数が降順だと一番小さくなります。だから、そのような数だけまず考えます。ここで、指数を並べると、12は[2, 1]と書けて、50は[1, 0, 2]と書けます。0を含まないで降順になっているものです。
そのような数はPriorityQueueで生成することができます。PriorityQueueはDynamicVectorで実現できますが、構造体をこの要素にするとき、以前はデコレーター@register_passableをつければいいと思っていたのですが、先月のアップデートでCollectionElementを直接実装すればよくなったようです。
from math import sqrt import sys #################### Item #################### var ps = DynamicVector[Int]() fn make_prime_table(N: Int): let a = Pointer[Bool].alloc(N+1) for i in range(N+1): a.store(i, True) var p: Int = 2 while p * p <= N: if not a.load(p): p += 1 continue for n in range(p*p, N+1, p): a.store(n, False) p += 1 for n in range(2, N+1): if a.load(n): ps.push_back(n) a.free() struct Item(CollectionElement): var n: Int var size: Int var es: Pointer[Int] fn __init__(inout self, n: Int, size: Int, es: Pointer[Int]): self.n = n self.size = size self.es = es fn __copyinit__(inout self, other: Self): self.n = other.n self.size = other.size self.es = Pointer[Int].alloc(other.size) for i in range(self.size): self.es[i] = other.es[i] fn __moveinit__(inout self, owned other: Self): self.n = other.n self.size = other.size self.es = other.es fn __del__(owned self): self.es.free() fn __lt__(self, other: Item) -> Bool: return self.n < other.n fn __le__(self, other: Item) -> Bool: return self.n <= other.n fn __gt__(self, other: Item) -> Bool: return self.n > other.n fn __ge__(self, other: Item) -> Bool: return self.n >= other.n fn extend(self) -> Item: let es = Pointer[Int].alloc(self.size+1) for i in range(self.size): es[i] = self.es[i] es[self.size] = 1 let n = self.n * ps[self.size] return Item(n, self.size+1, es) fn value(self) -> Int: var n: Int = 1 for i in range(self.size): let p = ps[i] let e = self.es[i] n *= p**e return n fn num_divs(self) -> Int: var n = 1 for i in range(self.size): n *= self.es[i] + 1 return n fn is_all_ones(self) -> Bool: for i in range(self.size): if self.es[i] != 1: return False return True fn is_descending(self) -> Bool: for i in range(1, self.size): if self.es[i-1] < self.es[i]: return False return True fn rfind(self, e: Int) -> Int: for i in range(self.size-1, -1, -1): if self.es[i] == e: return i return -1 fn append_one(self) -> Self: let n = self.n * ps[self.size] let es = Pointer[Int].alloc(self.size+1) for i in range(self.size): es[i] = self.es[i] es[self.size] = 1 return Item(n, self.size+1, es) fn mul(self, i: Int) -> Self: var item = self item.n *= ps[i] item.es[i] += 1 return item # 正規のItem（降順）の次 fn nexts(self) -> DynamicVector[Item]: var items = DynamicVector[Item]() items.push_back(self.mul(0)) if self.is_all_ones(): items.push_back(self.append_one()) for i in range(1, self.size): if self.es[i] == self.es[i-1] - 1: items.push_back(self.mul(i)) break elif self.es[i] != self.es[i-1]: break return items fn swap(self, i: Int) -> Self: let de = self.es[i] - self.es[i+1] var el = self let tmp = el.es[i] el.es[i] = el.es[i+1] el.es[i+1] = tmp el.n = el.n // ps[i]**de * ps[i+1]**de return el fn move_right(self) -> Self: let es = Pointer[Int].alloc(self.size+1) for i in range(self.size-1): es[i] = self.es[i] es[self.size-1] = 0 es[self.size] = self.es[self.size-1] let e = self.es[self.size-1] let n = self.n // ps[self.size-1]**e * ps[self.size]**e return Item(n, self.size+1, es) # 約数の数が同じ場合の次 fn nexts2(self) -> DynamicVector[Item]: var els = DynamicVector[Item]() let index = self.rfind(0) if index == -1: # all zero for i in range(1, self.size): if self.es[i-1] > self.es[i]: els.push_back(self.swap(i-1)) break else: # .., e, 0, f, .. -> .., 0, e, f, .. if index > 0 and self.es[index-1] != 0: els.push_back(self.swap(index-1)) els.push_back(self.move_right()) return els fn __str__(self) -> String: var s = String(self.n) + " [ " for i in range(self.size): s += str(self.es[i]) + " " s += "]" return s @staticmethod fn one() -> Item: let es = Pointer[Int].alloc(0) return Item(1, 0, es) #################### HeapQueue #################### struct HeapQueue: var a: DynamicVector[Item] fn __init__(inout self, a: DynamicVector[Item]): self.a = a fn __copyinit__(inout self, other: HeapQueue): self.a = DynamicVector[Item]() for i in range(other.a.size): self.a.push_back(other.a[i]) fn pop(inout self) -> Item: let top = self.a[0] self.a[0] = self.a.pop_back() let N = self.a.size var i = 0 while i < N-1: let j1 = i*2+1 let 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: Item): self.a.push_back(e) var i = self.a.size - 1 while i > 0: let 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): let tmp = self.a[i] self.a[i] = self.a[j] self.a[j] = tmp^ #################### HighCompositeGenerator #################### struct HighCompositeGenerator: var hq: HeapQueue var N: Int fn __init__(inout self, hq: HeapQueue, N: Int): self.hq = hq self.N = N fn __copyinit__(inout self, other: Self): self.hq = other.hq self.N = other.N fn next(inout self) -> Item: while True: let item = self.hq.pop() let items = item.nexts() for i in range(items.size): self.hq.push(items[i]) if item.num_divs() > self.N: return item @staticmethod def initialize(N: Int) -> HighCompositeGenerator: let es = Pointer[Int].alloc(1) es[0] = 1 let item = Item(2, 1, es) var a = DynamicVector[Item]() a.push_back(item) let hq = HeapQueue(a) return HighCompositeGenerator(hq, N) #################### HighCompositeGeneratorAll #################### struct HighCompositeGeneratorAll: var hcg: HighCompositeGenerator var hq: HeapQueue fn __init__(inout self, hcg: HighCompositeGenerator, hq: HeapQueue): self.hcg = hcg self.hq = hq fn next(inout self) -> Int: let item = self.hq.pop() if item.is_descending(): let item2 = self.hcg.next() self.hq.push(item2) let items = item.nexts2() for i in range(items.size): self.hq.push(items[i]) return item.n @staticmethod def initialize(N: Int) -> HighCompositeGeneratorAll: var hcg = HighCompositeGenerator.initialize(N) let item = hcg.next() var a = DynamicVector[Item]() a.push_back(item) let hq = HeapQueue(a) return HighCompositeGeneratorAll(hcg, hq) #################### process #################### fn is_triangle(n: Int) -> Bool: let m = sqrt(n*2) return m * (m + 1) // 2 == n fn f(N: Int) raises -> Int: var hq = HighCompositeGeneratorAll.initialize(N) while True: let n = hq.next() if is_triangle(n): return n fn main() raises: make_prime_table(100) let args = sys.argv() let N = atol(args[1]) print(f(N))
