Haskellで無限個の無限リストをソートされた形で結合する

CodeforcesProject Eulerの問題には、無限リストをうまく使うと綺麗に解くことができる問題がたくさんあります。 数列の性質から探索範囲の上界を決めて解を探索することが多いのですが、きちんとした根拠を持って上界を決めることができることは少なく、余裕を持って十分に広い範囲で計算して解を求める解法がよく取られます。

Haskellの特徴である遅延評価とその洗練された糖衣構文を用いると、無限リストを簡単に扱うことができます。 上界を適当に定める解法よりも、より宣言的で美しく、時に効率的なコードで同じ解を得ることができます。 しかし、無限リストをきちんと、それも無限個の無限リストをきちんと扱うとなると、意外と苦労します。

この記事では、無限個の無限リストをソートされた形で結合する方法について説明します。 一般的な無限リストではなく、条件はかなり絞っていてます (そうでないと原理的に解けません)。 まず、それぞれのリストはソートされていて、リストは無限個あるもののリストの要素間の大小は決まっているとします。 詳細は後々説明していきます。

いきなり無限個の無限リストを考えるのは難しいので、次の三通りを順番に考えて、簡単な問題から解いていくことにします。

  1. 2つの無限リスト
  2. 3つの無限リスト
  3. 無限個の無限リスト

マージする方法も和集合と共通集合があります。

  1. 和集合
  2. 共通集合

この掛け合わせの6通りについて、問題と解いていく形式で説明を進めていきます。 以下の6つがこのエントリーで扱う問題です。

【問題】

  • 三角数もしくは平方数である数のうち、小さい方から1000番目の数字を答えよ。
  • 三角数でかつ平方数である数のうち、小さい方から7番目の数字を答えよ。
  • 三角数または平方数または五角数である数のうち、小さい方から10万番目の数字を答えよ。
  • 三角数でかつ平方数でかつ六角数である数のうち、小さい方から5番目の数字を答えよ。
  • 正の整数のうち、ある奇素数pに対するp角数の要素であるような数の集合を考える。このような数の中で、小さい方から1万番目の数字を答えよ。
  • 21は三角数であり八角数であり21角数であるため、多角数として三通りの表現ができる。多角数としてちょうど五通りの表現ができる数のうち、小さい方から150番目の数字を答えよ。

これらの問題をご自分で解きたいと思われた方は、ここで一旦記事を読むのをやめて解いてみてください。 解きますか?解きませんか? もうこの下からは問題に対する解答が始まりますよ。

それでは、無限リストが織りなす宣言的な世界へようこそ。

2つの無限リストの和集合

【問題】三角数もしくは平方数である数のうち、小さい方から1000番目の数字を答えよ。

この記事ではghciで作業を進めていきます。 今回はpromptにモジュール名が出ていても特に嬉しいことはないので、隠しておきましょう。

 $ ghci
Prelude> :set prompt "> "
> :set prompt2 "| "

まず、平方数を定義してみます。

> let squares = [ n * n | n <- [1..] ]
> take 10 squares
[1,4,9,16,25,36,49,64,81,100]

次は三角数です。

> let triangulars = [ n * (n + 1) `div` 2 | n <- [1..] ]
> take 10 triangulars
[1,3,6,10,15,21,28,36,45,55]

これらを結合して小さい方から1000番目の数字を取り出しましょう。 単純な解法は、次のようなものが考えられるでしょう。

> import Data.List
> let n = 1000 in sort (union (take n squares) (take n triangulars)) !! (n - 1)
173166

答えが出ました。平方数または三角数の正の整数のうち、1000番目の数字は173166です。 1000番目を求めよという問題なので、それぞれの数列から1000個取ってきて和集合をとり、ソートして1000番目を返すという感じです。

しかし、このような問題に対しては、unionsortも使うべきではありません。 とても効率が悪く、100万番目を求めてくださいという問題だと、現実的な時間で解くことができません。

> let n = 1000000 in sort (union (take n squares) (take n triangulars)) !! (n - 1)
^CInterrupted.

しばらく待ちましたが、返ってきませんでした。

2つのリストが既にソートされていると分かっているのに、unionsortを使うは無駄な計算です。 結合してからソートするのではなく、まるでファスナーを上げるように、結合しながらソートされたリストを作ることができます。

> :{
| let union2 xs [] = xs
|     union2 [] ys = ys
|     union2 xxs@(x:xs) yss@(y:ys)
|       | x < y = x : union2 xs yss
|       | x == y = x : union2 xs ys
|       | otherwise = y : union2 xxs ys
| :}
> let n = 1000 in union2 squares triangulars !! (n - 1)
173166

同じ答えが出ました。 もちろん、この方法は最初に考えたunionしてsortする方法よりも圧倒的に高速です。

> let n = 100000 in union2 squares triangulars !! (n - 1)
1716013236
> let n = 1000000 in union2 squares triangulars !! (n - 1)
171575840736

10万番目だろうと100万番目だろうとすぐに答えが返ってきます。

このunion2関数は、ソートされた2つのリストをソートされた形でマージする関数です。 大事な関数なのでもう一度書いておきます。

union2 xs [] = xs
union2 [] ys = ys
union2 xxs@(x:xs) yss@(y:ys)
  | x < y = x : union2 xs yss
  | x == y = x : union2 xs ys
  | otherwise = y : union2 xxs ys

上記の問題の制約上、同じ数字は省いており、しかも入力のリストはどちらも真に増えていく (strictly increasing) 数列でなくてはなりません。 この条件を緩めて、単純に増えていく (同じ数字が続いても良い) 数列を受け付けて、かつ重複は省略しないならば、次のような実装になるでしょう。

union2' xs [] = xs
union2' [] ys = ys
union2' xxs@(x:xs) yss@(y:ys)
  | x <= y = x : union2' xs yss
  | otherwise = y : union2' xxs ys

この2つの関数の違いは、以下の結果で一目瞭然です。

> union2 [1..10] [1..10]
[1,2,3,4,5,6,7,8,9,10]
> union2' [1..10] [1..10]
[1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10]

いずれにせよ、入力がソートされているならば、効率よくソートされた結合リストを作ることができるということが分かります。 これら2つの関数はどちらが良いとか悪いとかではなく、入力に課している条件も結果も違いますので、問題によって使い分けることになります。 入力は同じ数字が続くかもしれない増加数列で、かつ重複しないようにマージする関数を書くことは、読者への課題とします (ヒント: dropWhile)。

2つの無限リストの共通集合

【問題】三角数でかつ平方数である数のうち、小さい方から7番目の数字を答えよ。

先ほどの問題では、1000番目と言われたときには三角数の1000番目までか、もしくは平方数の1000番目までかのいずれかに必ず含まれるということが分かっていました。 しかし、今回の場合、最初にいくつ取ればいいのか分かりません。

> intersect (take 100 triangulars) (take 100 squares)
[1,36,1225]

試しに100個ずつとってみて、共通するのは3つだけでした。 7番目と言われているので、簡単に見つかるのではと思うかもしれませんが、なかなか見つかりません。

> intersect (take 5000 triangulars) (take 5000 squares)
[1,36,1225,41616,1413721]
> intersect (take 10000 triangulars) (take 10000 squares)
[1,36,1225,41616,1413721,48024900]
> intersect (take 50000 triangulars) (take 50000 squares)
[1,36,1225,41616,1413721,48024900^CInterrupted.

すでに6個求まっているのであと少しなんですが、数を増やすとえらく時間がかかってしまいます。

ここで、triangularssquaresがソートされていることを思い出します。 しかも、どちらも真に増加していく数列です。 先ほどのunion2関数と同じ要領で、入力がソートされていると仮定したintersect2関数を定義してみます。

> :{
| let intersect2 _ [] = []
|     intersect2 [] _ = []
|     intersect2 xxs@(x:xs) yss@(y:ys)
|       | x < y = intersect2 xs yss
|       | x == y = x : intersect2 xs ys
|       | otherwise = intersect2 xxs ys
| :}
> take 7 $ intersect2 triangulars squares
[1,36,1225,41616,1413721,48024900,1631432881]

すぐに求まりました! 三角数でありかつ平方数である数の7番目の数は、1631432881です!

一見良さそうに見えますが、残念ながらこの方法でも20個求めようとなると時間がかかりすぎてしまいます。

> take 20 $ intersect2 triangulars squares
[1,36,1225,41616,1413721,48024900,1631432881,55420693056,1882672131025^CInterrupted.

この結果を見ると、だいたい34倍ずつになっていることが分かります。

> let xs = intersect2 triangulars squares in zipWith (\x y -> fromIntegral x / fromIntegral y) (tail xs) xs
[36.0,34.02777777777778,33.972244897959186,33.97061226451365,33.97056420609158,33.970562791385305,33.97056274974024,33.970562748514325^CInterrupted.

はて、何やら法則でもあるのかしらと思って、OEISで一般項を調べます。 OEISに「1,36,1225,41616」と入力すると、A001110が見つかります。 FORMULAのところを見ると、次のような漸化式が書かれています。

a(0) = 0, a(1) = 1; for n >= 2, a(n) = 34 * a(n-1) - a(n-2) + 2.

なるほど、漸化式に34という係数がありますね。 なぜこうなるのかはさておき、数学者が一生懸命研究して導き出した漸化式ですから、正しいことを期待してghciで定義して30番目まで求めてみましょう。

> let squareTriangulars = 0 : 1 : zipWith (\x y -> 34 * x - y + 2) (tail squareTriangulars) squareTriangulars
> take 30 squareTriangulars
[0,1,36,1225,41616,1413721,48024900,1631432881,55420693056,1882672131025,63955431761796,2172602007770041,73804512832419600,2507180834294496361,85170343853180456676,2893284510173841030625,98286503002057414584576,3338847817559778254844961,113422539294030403250144100,3853027488179473932250054441,130889512058808083293251706896,4446390382511295358038307980025,151046383493325234090009219613956,5131130648390546663702275158894481,174307395661785261331787346182798400,5921320321852308338617067495056251121,201150583547316698251648507485729739716,6833198520286915432217432187019754899225,232127599106207807997141045851185936833936,7885505171090778556470578126753302097454601]

一瞬で求まりました。 OEISの解答で答え合わせをして、あっていることを確認しましょう。 三角数でかつ平方数である数列を求める問題は、あるペル方程式の解に帰着することがWikipediaの記事やMathWorldの記事を見ると分かります。

3つの無限リストの和集合

【問題】三角数または平方数または五角数である数のうち、小さい方から10万番目の数字を答えよ。

まずは五角数を定義します。

> let pentagonals = [ n * (3 * n - 1) `div` 2 | n <- [1..] ]
> take 10 pentagonals
[1,5,12,22,35,51,70,92,117,145]

定義できました。

3つのソートされた (しかも真に増加していく) 数列 triangulars, squares そして pentagonals を小さい順にマージしていきます。 とりあえずどれも無限リストなので、引数が[]になることは考えないことにします。

> :{
| let union3 xxs@(x:_) yss@(y:_) zzs@(z:_)
|       = w : union3 (dropWhile (==w) xxs) (dropWhile (==w) yss) (dropWhile (==w) zzs)
|           where w = minimum [ x, y, z ]
| :}

<interactive>:96:5: Warning:
    Pattern match(es) are non-exhaustive
    In an equation for ‘union3’:
        Patterns not matched:
            [] _ _
            (_ : _) [] _
            (_ : _) (_ : _) []
> take 100 $ union3 triangulars squares pentagonals
[1,3,4,5,6,9,10,12,15,16,21,22,25,28,35,36,45,49,51,55,64,66,70,78,81,91,92,100,105,117,120,121,136,144,145,153,169,171,176,190,196,210,225,231,247,253,256,276,287,289,300,324,325,330,351,361,376,378,400,406,425,435,441,465,477,484,496,528,529,532,561,576,590,595,625,630,651,666,676,703,715,729,741,780,782,784,820,841,852,861,900,903,925,946,961,990,1001,1024,1035,1080]
> take 100 (union3 triangulars squares pentagonals) == take 100 (sort (take 100 triangulars `union` take 100 squares `union` take 100 pentagonals))
True

ちゃんと合ってそうですね。 10万番目の数を求めてみます。

> let n = 10000 in union3 triangulars squares pentagonals !! (n - 1)
9603153
> let n = 100000 in union3 triangulars squares pentagonals !! (n - 1)
958335849

求めることができました! 三角数または平方数または五角数である正の整数で、10万番目の数は958335849だということが分かりました。

3つの無限リストの共通集合

【問題】三角数でかつ平方数でかつ六角数である数のうち、小さい方から5番目の数字を答えよ。

まずは六角数を定義しましょう。

> let hexagonals = [ n * (2 * n - 1) | n <- [1..] ]
> take 10 hexagonals
[1,6,15,28,45,66,91,120,153,190]

3つの無限リストの共通集合です。 これらのリストは真に増加するリストなので、それぞれの最初の数を比較して小さいものは落としていけばよいことが分かります。

> :{
| let intersect3 xxs@(x:xs) yss@(y:ys) zzs@(z:zs)
|       | x < y || x < z = intersect3 xs yss zzs
|       | y < z || y < x = intersect3 xxs ys zzs
|       | z < x || z < y = intersect3 xxs yss zs
|       | otherwise = x : intersect3 xs ys zs
| :}

<interactive>:138:5: Warning:
    Pattern match(es) are non-exhaustive
    In an equation for ‘intersect3’:
        Patterns not matched:
            [] _ _
            (_ : _) [] _
            (_ : _) (_ : _) []

引数が[]であるときに対応できていないという警告が出ていますが、とりあえず今は無限リストに対して使うので無視します。 三角数でかつ平方数でかつ六角数であるような数は、次のようになります。

> take 5 $ intersect3 triangulars squares hexagonals
[1,1225,1413721,1631432881,1882672131025]

従って、求めるべき5番目の数は1882672131025であることが分かりました。 OEISに「1,1225,1413721,1631432881,1882672131025」と入力すると、A046177が見つかり、確かに答えがあっていることを確認できます。 OEISに書かれている漸化式をコードに落として30番目まで求めるのは、読者への課題とします。

実はこの問題はintersect3を使わなくても解けてしまいます。 全ての六角数は三角数でもあるからです。

> take 1000 (intersect2 triangulars hexagonals) == take 1000 hexagonals
True

従って、intersect2を使えば解けてしまいます。

> take 5 $ intersect2 squares hexagonals
[1,1225,1413721,1631432881,1882672131025]

同じ答になりましたね。 なんじゃそれと思われるかもしれませんが、問題を解く前にきちんと数学的な性質を調べて無駄な計算をしないことはとても大切なことなのです。

ここまで理解できた方はProject Euler 45を簡単に解くことができるでしょう。

無限個の無限リストの和集合

【問題】正の整数のうち、ある奇素数pに対するp角数の要素であるような数の集合を考える。このような数の中で、小さい方から1万番目の数字を答えよ。

まずは、Wikipediaの記事から多角数表を引っ張ってきます。

k n=1 n=2 n=3 n=4 n=5 n=6 n=7
3 1 3 6 10 15 21 28
4 1 4 9 16 25 36 49
5 1 5 12 22 35 51 70
6 1 6 15 28 45 66 91
7 1 7 18 34 55 81 112
8 1 8 21 40 65 96 133
9 1 9 24 46 75 111 154
10 1 10 27 52 85 126 175
11 1 11 30 58 95 141 196
12 1 12 33 64 105 156 217
13 1 13 36 70 115 171 238

素数pp角数と言われているので、該当する行のみ抽出します。

k n=1 n=2 n=3 n=4 n=5 n=6 n=7
3 1 3 6 10 15 21 28
5 1 5 12 22 35 51 70
7 1 7 18 34 55 81 112
11 1 11 30 58 95 141 196
13 1 13 36 70 115 171 238

この表に含まれている数を、小さい方から並べていきます。

1, 3, 5, 6, 7, 10, 11, 12, 13, 15 ...

すぐに分かることは、この数列は素数を全て含んでいるということです。 素数以外にも 1, 6, 10, 12, 15 のような数が含まれています。

素数は無限個あります。 ある素数pに対して、p角数は無限リストです。 すなわち、無限個の無限リストを、結合して小さい方から並べなさいという問題なのです。 無限個の無限リスト… この問題は解くことができるのでしょうか?

これまではせいぜい2つや3つの無限リストをマージしてきました。 しかし、同じ考え方では今回の問題は解けそうにありません。

> let union' xxs@(x:_) yss@(y:_) zzs@(z:_) wws@(w:_) uus@(u:_) ...

いえいえ、もちろん引数はリストのリストですよね。

> let union' xss = let w = minimum (map head xss) in ...

map head xssは無限リストなので、minimumを取ることができません。

私たちは、「奇素数角数表」を眺めながら、小さい方から順番に数字を言うことができます。 それは、この表が右に行くほど、下に行くほど大きくなるという性質があるからです。 改めて、表を見ながら私たちが小さい方から数字を言っていく時に何を考えているのか調べてみましょう。

まず、1はどの行にも一列目にあります。 次に、一列目を除いたら右に行くほど下に行くほど大きくなるので、次の数字は二列目以降の左上にある3ということが分かります。

n=1 n=2 n=3 n=4 n=5 n=6 n=7
1 3 6 10 15 21 28
1 5 12 22 35 51 70
1 7 18 34 55 81 112
1 11 30 58 95 141 196
1 13 36 70 115 171 238

次に比較するのは6, 5です。 小さい5を取ります。

n=1 n=2 n=3 n=4 n=5 n=6 n=7
1 3 6 10 15 21 28
1 5 12 22 35 51 70
1 7 18 34 55 81 112
1 11 30 58 95 141 196
1 13 36 70 115 171 238

次に、6, 12, 7で比較して、6を選びます。

n=1 n=2 n=3 n=4 n=5 n=6 n=7
1 3 6 10 15 21 28
1 5 12 22 35 51 70
1 7 18 34 55 81 112
1 11 30 58 95 141 196
1 13 36 70 115 171 238

10, 12, 7で比較して、7を選びます。

n=1 n=2 n=3 n=4 n=5 n=6 n=7
1 3 6 10 15 21 28
1 5 12 22 35 51 70
1 7 18 34 55 81 112
1 11 30 58 95 141 196
1 13 36 70 115 171 238

ようやく私たちがどのように数を列挙しているか分かってきたことでしょう。 比較すべき数はどんどん増えていきます。 赤くなっている数字よりも右や下は考える必要はありません。 7を選んだ時、1811が次に比較すべき数字の対象に追加されます。 次に101112と選ぶと、15, 22, 18, 30, 13が比較対象となります。

ルールが分かってきたのでコードに落とし込める気がしてきました。 頑張って実装していきましょう。

まずは素数を定義します。 適当なパッケージを使っても良いですが、素数のリストを定義するのは2行書いてあげれば事足ります。 「素数のリスト」は「2と素数である奇数」で、「ある数が素数である」とは「その数の平方根以下の素数でその数を割り切れない」ということです。

> :{
| let primes = 2 : filter isPrime [3,5..]
|     isPrime n = n > 1 && foldr (\p r -> p * p > n || n `mod` p /= 0 && r) True primes
| :}
> take 100 primes
[2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97,101,103,107,109,113,127,131,137,139,149,151,157,163,167,173,179,181,191,193,197,199,211,223,227,229,233,239,241,251,257,263,269,271,277,281,283,293,307,311,313,317,331,337,347,349,353,359,367,373,379,383,389,397,401,409,419,421,431,433,439,443,449,457,461,463,467,479,487,491,499,503,509,521,523,541]

きちんと素数になっていますね。 大丈夫そうです。

次に、多角数を定義します。 三角数自然数の和ですし、平方数は奇数の和です。 一般に、等差数列を足していくと多角数になります。

> let polygonals k = scanl1 (+) [1,k-1..]
> take 10 $ polygonals 3
[1,3,6,10,15,21,28,36,45,55]
> take 10 $ polygonals 4
[1,4,9,16,25,36,49,64,81,100]
> take 10 $ polygonals 5
[1,5,12,22,35,51,70,92,117,145]
> take 10 $ polygonals 6
[1,6,15,28,45,66,91,120,153,190]
> take 1000 (polygonals 5) == take 1000 pentagonals
True
> take 1000 (polygonals 6) == take 1000 hexagonals
True

OKですね。

試しに、「奇素数角数表」を表示してみます。 奇素数ではない素数2だけなので、奇素数tail primesです。

> mapM_ print [ (p, take 10 $ polygonals p) | p <- take 10 (tail primes) ]
(3,[1,3,6,10,15,21,28,36,45,55])
(5,[1,5,12,22,35,51,70,92,117,145])
(7,[1,7,18,34,55,81,112,148,189,235])
(11,[1,11,30,58,95,141,196,260,333,415])
(13,[1,13,36,70,115,171,238,316,405,505])
(17,[1,17,48,94,155,231,322,428,549,685])
(19,[1,19,54,106,175,261,364,484,621,775])
(23,[1,23,66,130,215,321,448,596,765,955])
(29,[1,29,84,166,275,411,574,764,981,1225])
(31,[1,31,90,178,295,441,616,820,1053,1315])

この右にも下にも無限に続く数列たちを、頑張ってマージしていかなくてはいけません。

まず、作りたい和集合を取る関数の型は、

> :{
| let union' :: Ord a => [[a]] -> [a]

としましょう。 マージしながら、次にいくつ数を取って比較するかという数字を持ち越します。 最初は1つなので、適当にgo 1と書いてからgo関数を一所懸命実装します。

|     union' = go 1
|       where
|         go k xss = let (k', m, yss) = extractMin k xss in m : go k' yss

extractMinは、リストのリストからk個とってそれらの最初要素の中で最小のものを見つけつつ、その要素を落としたものを返します。 慎重に定義しましょう。

|         extractMin k xss = (max k (l + 1), m, yss)
|           where
|             m = minimum $ map head $ take k xss
|             (l, yss) = extract m 1 xss

mは、リストのリストxssからk個とったそれらの最初の要素たちの最小です。 extractはリストのリストxssから初めて見つけたmを抜き出す関数です。 lmを何番目のリストで見つけたかという数字が返ってきます。

|         extract m l (xxs@(x:xs):yss)
|           | x == m = (l, xs:yss)
|           | otherwise = fmap (xxs:) $ extract m (l + 1) yss
| :}

<interactive>:25:9: Warning:
    Pattern match(es) are non-exhaustive
    In an equation for ‘extract’:
        Patterns not matched:
            _ _ []
            _ _ ([] : _)

最初のリストの最初の要素がmであれば、それを取ってxs:yssを返します。 extract関数で引数が[][]:_であるケースに対応できていないと怒られていますが、とりあえず今考えているのは無限リストなので放っておきます。

ようやく定義できたunion'関数をさっそく使ってみましょう。

> take 30 $ union' [ polygonals p | p <- tail primes ]
[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]

おっと、そうでした。 どの多角数も最初は1なので、これを除外しないと「右に行くほど下に行くほど大きく」という条件が達成されません。 それぞれの多角数の最初は無視することにします。

> take 30 $ union' [ tail (polygonals p) | p <- tail primes ]
[3,5,6,7,10,11,12,13,15,17,18,19,21,22,23,28,29,30,31,34,35,36,36,37,41,43,45,47,48,51]

一見良さそうに見えますが、よく見ると36が被ってしまっていますね。 36三角数でありかつ十三角数でもあります。 実は上のextract関数では、数が重複した時にそれを一つだけにしていません。 ソートされたリストの重複を除くのは、map head . groupを使うのがお手軽です。 あと最初の1を追加しておきます。

> let primePolygonals = 1 : map head (group (union' [ tail (polygonals p) | p <- tail primes ]))
> take 30 primePolygonals
[1,3,5,6,7,10,11,12,13,15,17,18,19,21,22,23,28,29,30,31,34,35,36,37,41,43,45,47,48,51]

やったー! できました! 重複もしていません。 この結果を表とじっと見比べたり、

k n=1 n=2 n=3 n=4 n=5 n=6 n=7
3 1 3 6 10 15 21 28
5 1 5 12 22 35 51 70
7 1 7 18 34 55 81 112
11 1 11 30 58 95 141 196
13 1 13 36 70 115 171 238
17 1 17 48 94 155 231 322
19 1 19 54 106 175 261 364
23 1 23 66 130 215 321 448

あるいは

> take 100 primePolygonals == take 100 (sort (foldl1 union [ take 100 (polygonals p) | p <- take 100 (tail primes) ]))
True

愚直に求めたものと合致することを確認して正しいことを確認します。 最後に問題に対する答を出します。

> let n = 10000 in primePolygonals !! (n - 1)
40641

出ました! 奇素数の多角数の中で小さい方から数えて1万番目の数は40641ということが分かりました! 無限個の無限リストをマージして小さい方から列挙することができました!

上記のコードをファイルにして実行できる形にしておきます。

import Data.List (group)

main :: IO ()
main = do
  print $ take 100 primePolygonals
  let n = 10000
  print $ primePolygonals !! (n - 1)

primes :: Integral a => [a]
primes = 2 : filter isPrime [3,5..]

isPrime :: Integral a => a -> Bool
isPrime n = n > 1 && foldr (\p r -> p * p > n || n `mod` p /= 0 && r) True primes

polygonals :: Integral a => a -> [a]
polygonals k = scanl1 (+) [1,k-1..]

primePolygonals :: Integral a => [a]
primePolygonals = 1 : map head (group (union' [ tail (polygonals p) | p <- tail primes ]))

union' :: Ord a => [[a]] -> [a]
union' = go 1
  where
    go k xss = let (k', m, yss) = extractMin k xss in m : go k' yss
    extractMin k xss = (max k (l + 1), m, yss)
      where
        m = minimum $ map head $ take k xss
        (l, yss) = extract m 1 xss
    extract m l (xxs@(x:xs):yss)
      | x == m = (l, xs:yss)
      | otherwise = fmap (xxs:) $ extract m (l + 1) yss

意外とシンプルでしょう? union'関数の引数が有限長のリストであったり有限個だったりするとうまく動かないので、これらに対応させるのは読者への課題とします (ヒント: concatMap (take 1))。

この無限個の無限リストを結合して小さい方から列挙する関数union'を用いると、Project Euler 126をきれいに解くことができます。 実はこの問題は、このエントリーを書き始めたきっかけの問題です。 Project Euler 126を解いている過程で、無限個のリストを処理する必要があり、書いてみたら意外とちゃんと動いたので勢いでこのエントリーも書いたというわけです。 無限リストを自在に操れるようになると、「この問題にはいくら以下でおそらく答えが見つかるだろうからこのくらいの長さの配列を作ってその中で解を探す」みたいな雑な見積もりをせずとも、宣言的ですっきりとしたコードを書くことができます。

無限個の無限リストの中のいくつかに含まれる集合

【問題】21は三角数であり八角数であり21角数であるため、多角数として三通りの表現ができる。多角数としてちょうど五通りの表現ができる数のうち、小さい方から150番目の数字を答えよ。

今回の問題では、「多角数として五通りの表現ができる数」となっています。 全ての多角数に含まれる数は1しかありません。 なぜなら全ての多角数の最初の要素は1である上に、任意の2以上の数kk角数に含まれないからです。 「五つの多角数の数列に含まれている」と表現することで、問題が数学的に意味を持つのです。

実は、先ほどのunion'関数は重複を除外しないので、これを使うと問題を簡単に解くことができます。

main :: IO ()
main = do
  print $ take 100 polygonalsIntersection
  let n = 150
  print $ polygonalsIntersection !! (n - 1)

polygonalsIntersection :: Integral a => [a]
polygonalsIntersection = map head (filter ((==5) . length) (group (union' [ tail (polygonals p) | p <- [3..] ])))

groupしてその長さが5となるものでfilterするだけです。 実行すると次のようになります。

[225,231,276,325,435,441,540,595,616,651,820,861,936,946,1035,1089,1128,1275,1288,1296,1326,1431,1521,1596,1716,1881,1926,1936,2025,2080,2145,2205,2211,2296,2380,2415,2485,2541,2640,2745,2835,2871,2916,2925,2961,3136,3160,3186,3201,3312,3381,3430,3486,3537,3655,3745,3816,3825,3861,4096,4200,4221,4225,4236,4356,4371,4465,4564,4576,4641,4656,4845,4941,4950,4992,5005,5076,5151,5160,5265,5292,5356,5481,5551,5625,5656,5685,5776,5886,5929,5995,6076,6084,6105,6111,6201,6345,6441,6480,6501]
10206

従って、答は10206です! それぞれの数字がどの多角数に属するか求めるのも難しくありません。

> let elem' x xs = x == head (dropWhile (<x) xs)
> mapM_ print [ [ (p, q) | q <- [3..p], elem' p (polygonals q) ] | p <- polygonalsIntersection ]
[(225,4),(225,8),(225,24),(225,76),(225,225)]
[(231,3),(231,6),(231,17),(231,78),(231,231)]
[(276,3),(276,6),(276,20),(276,93),(276,276)]
[(325,3),(325,6),(325,9),(325,34),(325,325)]
[(435,3),(435,6),(435,45),(435,146),(435,435)]
[(441,4),(441,14),(441,31),(441,148),(441,441)]
[(540,7),(540,10),(540,21),(540,181),(540,540)]
[(595,3),(595,15),(595,30),(595,61),(595,595)]
[(616,7),(616,13),(616,31),(616,104),(616,616)]
[(651,5),(651,9),(651,45),(651,218),(651,651)]
[(820,3),(820,20),(820,31),(820,138),(820,820)]
^CInterrupted.
>

225は15番目の平方数であり (1+3+5+..+29=15*15)、9番目の八角数であり (1+7+13+19+25+31+37+43+49)、5番目の24角数であり (1+23+45+67+89)、3番目の76角数であり (1+75+149)、そして最初の225角数です。

おわりに

CodeforcesProject Eulerには、2つ以上のパラメーターによって特徴づけられる数列を集計して答えを求める問題がたくさん出てきます。 いくつまでに答えが見つかるはずだから (この根拠は勘である場合もあったり数学的に証明できる上界だったりする・コンテスト中にあまりきちんと証明している暇はない) 固定長の配列を用意し、探索するパラメータの上限も適当なところで打ち切り、一気に二重ループを回してしまってから配列から答を得るという方法がよく取られます。

このエントリーでは、多角数を題材にとって、無限個の無限リストを扱う方法を紹介してみました。 多角数でなくても「右に行くほど下に行くほど大きくなる」というような構造がある場合は、このエントリーの関数をそのまま使うことができますし、そうでなくてもチェックする数を持ち越していくやり方は一般的に使えるテクニックでしょう。 無限リストをきちんと扱えるようになると、解を探索する上界を意識する必要がなくなりますし、コードも美しくなります。 また、無駄な計算をしないため、探索範囲を決めて計算する解法よりも速く計算できる可能性もあります。

大事なことは、データの構造や型の制約を考えて適切な関数を使うことです。 例えば簡単な例としては、10000未満の素数はいくつあるか?という問題には、filterではなくてtakeWhileを使うでしょう。 これは素数リストがどんどん大きくなるという性質を使っています。 既にソートされているリストの重複を除きなさいと言われたら、nubではなくてgroupを使うでしょう。 リストがソートされていないとしても、順序がつく型 (Ordinstance) ならば存在する要素をSetで持ち回すとnubよりも計算オーダーの低い関数を書くことができます。 つまり、型に制約を課していくと、その型の構造をうまく用いたデータ構造を使って計算量を落とすことができる可能性があるということです。

Haskellは遅延評価の評価戦略を取っており、また豊富な糖衣構文により、無限リストを簡単に扱うことができます。 Project Eulerのように数学的な問題を解く上で、無限リストをどれだけうまく使いこなせるかというのは重要なことです。 無限個の無限リストをきちんと扱うのは大変なことです。 目の前にあるデータの数学的な性質をしっかりと見抜き、効率のよいアルゴリズムを考えて、問題を解くことを楽しんでください。