我们要解决的第一个问题是知道我们可以完全填充多少轮，以及我们还需要分发多少巴斯。我们可以将其表示为sum表达式：

我们可以用这个来计算Bath的 k 的数量，以获得具有 M 的完整行的数量：

如果我们只考虑 * 完整 * 轮，第一个人将拥有 M×（M+1）/2，第二个人拥有 M×（M+1）/2+M，第 j 个人拥有 M×（M+1）/2+j×M。
有可能还有巴斯的剩余，实际上，在 M 轮之后，还有 r = K-M ×n×（m+n）/2 剩余，这些然后在最后一轮中分配。这里，min（M+1，r），第二个 min（M+2，max（r-M-1，0）），第三个 min（M+3，max（r-2×M-3，0）），等等。
如果我们这样假设所有的算术运算都在 O（1） 内运行，那么下面的算法在 O（n） 内运行：

distribute :: Int -> Int -> [Int]
distribute n k = zipWith (+) (take n [n0, n0 + m ..]) (final (m + 1) r)
  where
    m = (floor (sqrt (fromIntegral (8 * k * n + n * n * n * n) :: Double)) - n * n) `div` (2 * n)
    r = k - (m * n * (m + n) `div` 2)
    n0 = m * (m + 1) `div` 2

final :: Int -> Int -> [Int]
final i j = k : final (i + 1) (j - k)
  where
    k = min i j

我们可以通过列举可能的浴室来检查结果，例如有四个人：

Main> mapM_ print (map (distribute 4) [0 .. 50])
[0,0,0,0]
[1,0,0,0]
[1,1,0,0]
[1,2,0,0]
[1,2,1,0]
[1,2,2,0]
[1,2,3,0]
[1,2,3,1]
[1,2,3,2]
[1,2,3,3]
[1,2,3,4]
[2,2,3,4]
[3,2,3,4]
[3,3,3,4]
[3,4,3,4]
[3,5,3,4]
[3,5,4,4]
[3,5,5,4]
[3,5,6,4]
[3,5,7,4]
[3,5,7,5]
[3,5,7,6]
[3,5,7,7]
[3,5,7,8]
[3,5,7,9]
[4,5,7,9]
[5,5,7,9]
[6,5,7,9]
[6,6,7,9]
[6,7,7,9]
[6,8,7,9]
[6,9,7,9]
[6,9,8,9]
[6,9,9,9]
[6,9,10,9]
[6,9,11,9]
[6,9,12,9]
[6,9,12,10]
[6,9,12,11]
[6,9,12,12]
[6,9,12,13]
[6,9,12,14]
[6,9,12,15]
[7,9,12,15]
[8,9,12,15]
[9,9,12,15]
[10,9,12,15]
[10,10,12,15]
[10,11,12,15]
[10,12,12,15]
[10,13,12,15]