如何用Haskell实现O(n)时间复杂度的泰铢分配算法？

问题分析 #

给定X泰铢和N个人，按规则分配：

• 第t轮给第i个人分配 i + t - 1 泰铢（第1轮i=1得1、i=2得2…第2轮i=1得2、i=2得3，以此类推）
• 持续分配直到资金耗尽，输出每个人最终获得的金额数组

直接模拟每一轮分配的时间复杂度是O(k*n)（k为轮数），当X很大时效率低下。我们可以通过数学推导找到规律，将时间复杂度优化到O(n)。

核心思路 #

1. 计算完整轮数m：找到最大的m，使得前m轮分配的总金额不超过X。前m轮总金额公式为：
   $$\text{total}_m = \frac{n \times m \times (m + n)}{2}$$
   通过浮点计算后验证的方式快速找到符合条件的m。

2. 计算基础金额：每个人在m轮完整分配中获得的金额为：
   $$\text{base}_i = \frac{m \times (2i + m - 1)}{2}$$
   这是第i个人在每一轮分配的金额之和（第t轮得i+t-1，求和从t=1到m）。

3. 处理剩余资金：计算剩余资金r = X - total_m，在第m+1轮中依次给每个人分配i + m泰铢（直到r耗尽），不足的部分取剩余金额，后续的人分不到钱。

代码实现 #

import Data.List (unfoldr)

-- 找到最大的完整轮数m
findM :: Integer -> Integer -> Integer
findM n x = let maxM0 = floor $ (sqrt (fromIntegral (n*n + 8*x)) - fromIntegral n) / 2
                adjust m
                  | n*m*(m + n) > 2*x = adjust (m - 1)
                  | n*(m+1)*(m+1 + n) <= 2*x = adjust (m + 1)
                  | otherwise = m
            in adjust maxM0

-- 生成每个人在m轮完整分配中的金额
baseAmounts :: Integer -> Integer -> [Integer]
baseAmounts n m = [ m*(2*i + m - 1) `div` 2 | i <- [1..n] ]

-- 生成剩余资金的分配数组
remainingAmounts :: Integer -> Integer -> Integer -> [Integer]
remainingAmounts n r m = take n $ unfoldr genNext (1, r)
  where genNext (i, rem)
          | rem <= 0 || i > n = Nothing
          | otherwise = let amt = min (i + m) rem
                        in Just (amt, (i + 1, rem - amt))

-- 主分配函数
allocate :: Integer -> Integer -> [Integer]
allocate x n
  | x <= 0 = replicate n 0
  | n <= 0 = []
  | otherwise = let m = findM n x
                    totalBase = n*m*(m + n) `div` 2
                    r = x - totalBase
                    base = baseAmounts n m
                    rem = remainingAmounts n r m
                in zipWith (+) base rem

示例验证 #

输入allocate 21 5，输出[3,5,4,4,5]，与题目示例一致：

• 完整轮数m=1，基础金额为[1,2,3,4,5]
• 剩余资金r=6，第2轮分配[2,3,1,0,0]
• 两者相加得到最终结果

内容的提问来源于stack exchange，提问作者willeM_ Van Onsem