Обратный поиск со списком Monad в Haskell

Если у нас есть сумма S и смещение k, мы можем, таким образом, выяснить, существует ли l, для которого сумма верна:

2S = l(l+1)-k(k-1)

0=l²+l-2S-k(k-1)

таким образом, мы можем решить это уравнение с помощью:

l=(-1 + (1+8S+4k(k-1)))/2

Если это целое число, то последовательность существует. Например, для S=9 и k=4 мы получаем:

l = (-1 + (1+72+48))/2 = (-1 + 11)/2 = 10/2 = 5.

Мы можем использовать некоторые функции, такие как вавилонский метод [wiki] для быстрого вычисления целых квадратных корней:

squareRoot :: Integral t => t -> t
squareRoot n 
   | n > 0    = babylon n
   | n == 0   = 0
   | n < 0    = error "Negative input"
   where
   babylon a   | a > b = babylon b
               | otherwise = a
      where b  = quot (a + quot n a) 2

Мы можем проверить, действительно ли найденный корень является точным квадратным корнем, возведя корень в квадрат, и посмотреть, получим ли мы исходный ввод.

Итак, теперь, когда у нас есть это, мы можем перебрать нижнюю границу последовательности и найти верхнюю границу. Если она существует, мы возвращаем последовательность, в противном случае пробуем следующую:

decompose :: Int -> [[Int]]
decompose s = [ [k .. div (sq-1) 2 ]
              | k <- [1 .. s]
              , let r = 1+8*s+4*k*(k-1)
              , let sq = squareRoot r
              , r == sq*sq
              ]

Таким образом, мы можем, например, получить элементы с помощью:

Prelude> decompose 1
[[1]]
Prelude> decompose 2
[[2]]
Prelude> decompose 3
[[1,2],[3]]
Prelude> decompose 3
[[1,2],[3]]
Prelude> decompose 1
[[1]]
Prelude> decompose 2
[[2]]
Prelude> decompose 3
[[1,2],[3]]
Prelude> decompose 4
[[4]]
Prelude> decompose 5
[[2,3],[5]]
Prelude> decompose 6
[[1,2,3],[6]]
Prelude> decompose 7
[[3,4],[7]]
Prelude> decompose 8
[[8]]
Prelude> decompose 9
[[2,3,4],[4,5],[9]]
Prelude> decompose 10
[[1,2,3,4],[10]]
Prelude> decompose 11
[[5,6],[11]]

Мы можем дополнительно ограничить диапазоны, например указать, что k‹l, с помощью:

decompose :: Int -> [[Int]]
decompose s = [ [k .. l ]
              | k <- [1 .. div s 2 ]
              , let r = 1+8*s+4*k*(k-1)
              , let sq = squareRoot r
              , r == sq*sq
              , let l = div (sq-1) 2
              , k < l
              ]

Затем это дает нам:

Prelude> decompose 1
[]
Prelude> decompose 2
[]
Prelude> decompose 3
[[1,2]]
Prelude> decompose 4
[]
Prelude> decompose 5
[[2,3]]
Prelude> decompose 6
[[1,2,3]]
Prelude> decompose 7
[[3,4]]
Prelude> decompose 8
[]
Prelude> decompose 9
[[2,3,4],[4,5]]
Prelude> decompose 10
[[1,2,3,4]]
Prelude> decompose 11
[[5,6]]