Python – 统计所有从1到N中的K个数的组合，它们的和等于N。

huangapple
go评论104阅读模式
英文:

Python - Count all combinations of K numbers from 1-N whose sum is equal to N

问题

如何计算所有从1到n中选取k个数字的组合，使它们的总和等于n？例如，对于n = 10，k = 3，我们有(1, 2, 7)，(1, 3, 6)，(1, 4, 5)，(2, 3, 5)。

我尝试过使用itertools.combination，但对于大数字而言，它增长得非常快。

英文:

How do i count all combinations of k numbers from 1-n whose sum is equal to n? Like for n = 10, k = 3, we have (1, 2, 7), (1, 3, 6), (1, 4, 5), (2, 3, 5)

I've tried using itertools.combination but it grows really fast for big numbers

答案1

得分: 2

我们可以用 k 个不同的正整数制造出的最小数是 choose(k+1, 2)。

令 r(n,k) = n - choose(k+1, 2)。

那么从 k 个不同的整数中形成 n 的方式数量等于将 k 个非负、不一定不同的整数相加以得到 r(n,k) 的方式数量。思路是从 1, 2, 3, ..., k 开始，然后以非递减的方式分配 r(n,k) 到这些起始整数上。

例如，对于 10, 3：

1 + 2 + 3 = choose(4,2) = 6，所以 r(10,3) = 10-6 = 4。
4 = 0+0+4, 0+1+3, 0+2+2, 1+1+2
(1,2,3) + (0,0,4) = (1,2,7)
(1,2,3) + (0,1,3) = (1,3,6)
(1,2,3) + (0,2,2) = (1,4,5)
(1,2,3) + (1,1,2) = (2,3,5)

所以我们将问题简化为计算 k 个非负整数相加得到 r(n,k) 的方式数量。在[这里][1]有答案。

Ruby 代码（包括实用函数）：

  1. def choose(n, k)
  2.   return 0 if k > n
  3.   result = 1
  4.   1.upto(k) do |d|
  5.     result *= n
  6.     result /= d
  7.     n -= 1
  8.   end
  9.   return result
  10. end
  12. def count_partitions_nozeroes(n, k, cache = {})
  13.   return 1 if k==0 && n==0
  14.   return 0 if n <= 0 || k <= 0
  16.   # 检查结果是否已被记忆
  17.   if cache.key?([n, k])
  18.     return cache[[n, k]]
  19.   end
  21.   # 计算结果
  22.   result = count_partitions_nozeroes(n-1, k-1, cache) + count_partitions_nozeroes(n-k, k, cache)
  24.   # 记忆结果以备将来使用
  25.   cache[[n, k]] = result
  26.   return result
  27. end
  29. def count_partitions_zeros(n,k)
  30.   return count_partitions_nozeroes(n+k, k)
  31. end
  33. def solve(n,k)
  34.   r = n - choose(k+1,2)
  35.   return count_partitions_zeros(r,k)
  36. end

示例结果：

  1. > solve(10,3)
  2. => 4
  4. > solve(200,10)
  5. => 98762607
  7. > solve(2000,10)
  8. => 343161146717017732
  10. > solve(2000,100) # 没有解决方案是正确的
  11. => 0
  13. > solve(2000,40)
  14. => 2470516759655914864269838818691
  16. > solve(5000,50)
  17. => 961911722856534054414857561149346788190620561928079
  19. > solve(9000,100)
  20. => 74438274524772625088229884845232647085568457172246625852148213

这是一个更简单的 Ruby 版本，避免了递归（其他方法不变）。它给出与上述相同的结果。下面展示了一些更大数字的结果。这个版本的复杂度是 O(n*r)。

  1. def count_partitions_nozeroes(n, k)
  2.   n_to_k_to_count = Hash.new{|h, n| h[n] = Hash.new{|h2, k| h2[k] = 0}}
  3.   n_to_k_to_count[n][k] = 1
  4.   
  5.   (n).downto(1) do |cur_n|
  6.     n_to_k_to_count.delete(cur_n + 1) # 删除旧键以节省空间
  7.     n_to_k_to_count[cur_n].keys.each do |cur_k|
  8.       n_to_k_to_count[cur_n - 1][cur_k - 1] += n_to_k_to_count[cur_n][cur_k] if cur_n >= 1 && cur_k >= 1
  9.       n_to_k_to_count[cur_n - cur_k][cur_k] += n_to_k_to_count[cur_n][cur_k] if cur_n >= cur_k && cur_k >= 0
  10.     end
  11.   end
  12.   return n_to_k_to_count[0][0] 
  13. end

示例结果：

  1. > solve(10_000, 100)
  2. => 274235043379646744332574760930015102932669961381003514201948469288939
  4. > solve(20_000, 100)
  5. => 7299696028160228272878582999080106323327610318395689691894033570930310212378988634117070675146218304092757
  7. > solve(30_000, 100)
  8. => 272832080760303721646457320315409638838332197621252917061852201523368622283328266190355855228845140740972789576932357443034296
  10. > solve(40_000, 200)
  11. => 1207940070190155086319681977786735094825631330761751426889808559216057614938892266960158470822904722575922933920904751545295375665942760497367
  13. > solve(100_000, 200)
  14. => 13051215883535384859396062192804954511590479767894013629996324213956689010966899432038449004533035681835942448619230013858515264041486939129111486281204426757510182253404556858519289275662797170197384965998425620735381780708992863774464769
  16. > solve(1_000_000, 200) # 变得非常慢；3.5 分钟
  17. => 428880856178598710720148624933560494061607079247573557573778067722670591454531582929217788942407876811003263888596981076595546473767426764847052870957098719920895206333233661830556744660481006393060648337767876434226805997102371290790505388472758064159747958795845134023811256732973394383039538732268993828238034324648751357082834429815006950891214256221354725682849015159958577756592134668188434645414960901194459625871943042
  19. <details>
  20. <summary>英文:</summary>
  22. The smallest number we can make with k distinct positive integers is choose(k+1, 2).
  24. Let r(n,k) = n - choose(k+1, 2).
  26. Then the count of ways of forming n from k distinct integers is equal to the count of ways of summing k nonnegative non-necessarily-distinct integers to get r(n,k). The idea is we start with 1, 2, 3, ..., k, and then allocate r(n,k) to these starting integers in a nondecreasing manner.
  28. E.g., 10, 3:
  30.     1 + 2 + 3 = choose(4,2) = 6, so r(10,3) = 10-6 = 4.
  31.     4 = 0+0+4, 0+1+3, 0+2+2, 1+1+2
  32.     (1,2,3) + (0,0,4) = (1,2,7)
  33.     (1,2,3) + (0,1,3) = (1,3,6)
  34.     (1,2,3) + (0,2,2) = (1,4,5)
  35.     (1,2,3) + (1,1,2) = (2,3,5)
  37. So we&#39;ve reduce the problem to counting the number of ways of summing k nonnegative integers to get r(n,k). Answered [here][1]
  39. Ruby code (including util functions):
  41.     def choose(n, k)
  42.       return 0 if k &gt; n
  43.       result = 1
  44.       1.upto(k) do |d|
  45.         result *= n
  46.         result /= d
  47.         n -= 1
  48.       end
  49.       return result
  50.     end
  51.     
  52.     def count_partitions_nozeroes(n, k, cache = {})
  53.       return 1 if k==0 &amp;&amp; n==0
  54.       return 0 if n &lt;= 0 || k &lt;= 0
  55.     
  56.       # Check if the result is already memoized
  57.       if cache.key?([n, k])
  58.         return cache[[n, k]]
  59.       end
  60.     
  61.       # Calculate the result
  62.       result = count_partitions_nozeroes(n-1, k-1, cache) + count_partitions_nozeroes(n-k, k, cache)
  63.     
  64.       # Memoize the result for future use
  65.       cache[[n, k]] = result
  66.       return result
  67.     end
  68.     
  69.     def count_partitions_zeros(n,k)
  70.       return count_partitions_nozeroes(n+k, k)
  71.     end
  72.     
  73.     def solve(n,k)
  74.       r = n - choose(k+1,2)
  75.       return count_partitions_zeros(r,k)
  76.     end
  78. Sample results
  80.     &gt; solve(10,3)
  81.     =&gt; 4
  82.     
  83.     &gt; solve(200,10)
  84.     =&gt; 98762607
  85.     
  86.     &gt; solve(2000,10)
  87.     =&gt; 343161146717017732
  88.     
  89.     &gt; solve(2000,100) # correct that there&#39;s no solution
  90.     =&gt; 0
  91.     
  92.     &gt; solve(2000,40)
  93.     =&gt; 2470516759655914864269838818691
  94.     
  95.     &gt; solve(5000,50)
  96.     =&gt; 961911722856534054414857561149346788190620561928079
  97.     
  98.     &gt; solve(9000,100)
  99.     =&gt; 74438274524772625088229884845232647085568457172246625852148213
  101. Here&#39;s a simpler Ruby version that avoids recursion (other methods unchanged). It gives the same results as above. A few results for larger numbers shown below. This version is O(n*r).
  103.     def count_partitions_nozeroes(n, k)
  104.       n_to_k_to_count = Hash.new{|h, n| h[n] = Hash.new{|h2, k| h2[k] = 0}}
  105.       n_to_k_to_count[n][k] = 1
  106.       
  107.       (n).downto(1) do |cur_n|
  108.         n_to_k_to_count.delete(cur_n + 1) # delete old keys to save space
  109.         n_to_k_to_count[cur_n].keys.each do |cur_k|
  110.           n_to_k_to_count[cur_n - 1][cur_k - 1] += n_to_k_to_count[cur_n][cur_k] if cur_n &gt;= 1 &amp;&amp; cur_k &gt;= 1
  111.           n_to_k_to_count[cur_n - cur_k][cur_k] += n_to_k_to_count[cur_n][cur_k] if cur_n &gt;= cur_k &amp;&amp; cur_k &gt;= 0
  112.         end
  113.       end
  114.       return n_to_k_to_count[0][0] 
  115.     end
  117. Sample results
  119.     &gt; solve(10_000, 100)
  120.     =&gt; 274235043379646744332574760930015102932669961381003514201948469288939
  121.     
  122.     &gt; solve(20_000, 100)
  123.     =&gt; 7299696028160228272878582999080106323327610318395689691894033570930310212378988634117070675146218304092757
  125.     &gt; solve(30_000, 100)
  126.     =&gt; 272832080760303721646457320315409638838332197621252917061852201523368622283328266190355855228845140740972789576932357443034296
  127.     
  128.     &gt; solve(40_000, 200)
  129.     =&gt; 1207940070190155086319681977786735094825631330761751426889808559216057614938892266960158470822904722575922933920904751545295375665942760497367
  131.     &gt; solve(100_000, 200)
  132.     =&gt; 13051215883535384859396062192804954511590479767894013629996324213956689010966899432038449004533035681835942448619230013858515264041486939129111486281204426757510182253404556858519289275662797170197384965998425620735381780708992863774464769
  134.     &gt; solve(1_000_000, 200) # getting painfully slow; 3.5 mins
  135.     =&gt; 42888085617859871072014862493356049406160707924757355757377806772267059145453158292921778894240787681100326388859698107659554647376742676484705287095709871992089520633323366183055674466048100639306064833776787643422680599710237129079050538847275806415974795879584513402381125673297339438303953873226899382823803432464875135708283442981500695089121425622135472568284901515995857775659213466818843464541496090119445962587194304280691087464026800781
  140.   [1]: https://math.stackexchange.com/questions/217597/number-of-ways-to-write-n-as-a-sum-of-k-nonnegative-integers
  142. </details>
  146. # 答案2
  147. **得分**: 1
  149. 使用带有缓存的递归方法可以在合理的时间内产生结果：
  151. ```python
  152. from functools import lru_cache
  154. @lru_cache(None)
  155. def countNK(n, k, t=None):
  156.     t = n if t is None else t                    # 跟踪目标部分和
  157.     if k == 0:  return int(t==0)                 # 空集可以求和为零
  158.     if t < 1: return 0                            # 仅有效目标
  159.     if k > n: return 0                            # 值不够
  160.     return countNK(n-1, k, t) + countNK(n-1, k-1, t-n)  # 合并计数
  • 递归需要以逐渐减小的n值来寻找目标
  • 它还需要在从目标中删除每个值后对较短的组合进行相同的操作
  • 这将多次合并相同的计算，因此需要缓存

输出：

  1. print(countNK(10,3))     # 4
  3. print(countNK(200,10))   # 98762607

如果您需要处理较大的n值（例如500+），您需要增加递归限制或将函数转换为迭代循环。

英文:

A recursive approach with caching can produce results in a reasonable time:

  1. from functools import lru_cache
  2. @lru_cache(None)
  3. def countNK(n,k,t=None):
  4. t = n if t is None else t                    # track target partial sum
  5. if k == 0:  return int(t==0)                 # empty set can sum to zero
  6. if t &lt; 1 :  return 0                         # valid target only
  7. if k &gt; n :  return 0                         # not enough values
  8. return countNK(n-1,k,t)+countNK(n-1,k-1,t-n) # combine counts
  • recursion needs to aim for a target using progressively smaller values of n
  • it also needs to do that for shorter combinations after removing each value from the target
  • this will combine the same calculations multiple times, hence the caching

...

output:

  1. print(countNK(10,3))     # 4
  2. print(countNK(200,10))   # 98762607

If you need to process large values of n (e.g. 500+), you'll either need to increase the recursion limit or convert the function to an iterative loop

答案3

得分: 1

以下是您提供的代码的中文翻译部分：

  1. def Kelly(n, k):
  2.     if k == 0:
  3.         return 1 if n == 0 else 0
  4.     @cache
  5.     def c(n, k):
  6.         if n < k * (k+1) // 2:
  7.             return 0
  8.         if k == 1:
  9.             return 1
  10.         n -= k
  11.         return c(n, k) + c(n, k-1)
  12.     return c(n, k)
  14. # 预先计算用于 "n < ..." 基本情况的界限
  15. def Kelly2(n, k):
  16.     if k == 0:
  17.         return 1 if n == 0 else 0
  18.     min_n_for_k = list(accumulate(range(k+1)))
  19.     @cache
  20.     def c(n, k):
  21.         if n < min_n_for_k[k]:
  22.             return 0
  23.         if k == 1:
  24.             return 1
  25.         n -= k
  26.         return c(n, k) + c(n, k-1)
  27.     return c(n, k)
  29. def Alain(n, k):
  30.     @lru_cache(None)
  31.     def countNK(n,k,t=None):
  32.         t = n if t is None else t                    # 跟踪目标部分和
  33.         if k == 0:  return int(t==0)                 # 空集可以总和为零
  34.         if t < 1 :  return 0                         # 仅有效目标
  35.         if k > n :  return 0                         # 不足够的值
  36.         return countNK(n-1,k,t)+countNK(n-1,k-1,t-n) # 合并计数
  37.     return countNK(n, k)
  39. def Dave_translated_by_Kelly(n, k):
  41.   def choose(n, k):
  42.     if k > n: return 0
  43.     result = 1
  44.     for d in range(1, k+1):
  45.       result *= n
  46.       result //= d
  47.       n -= 1
  48.     return result
  50.   def count_partitions_nozeroes(n, k, cache = {}):
  51.     if k==0 and n==0: return 1
  52.     if n <= 0 or k <= 0: return 0
  54.     # 检查结果是否已被记忆
  55.     if (n, k) in cache:
  56.       return cache[n, k]
  58.     # 计算结果
  59.     result = count_partitions_nozeroes(n-1, k-1, cache) + count_partitions_nozeroes(n-k, k, cache)
  61.     # 为将来的使用记忆结果
  62.     cache[n, k] = result
  63.     return result
  65.   def count_partitions_zeros(n,k):
  66.     return count_partitions_nozeroes(n+k, k)
  68.   def solve(n,k):
  69.     r = n - choose(k+1,2)
  70.     return count_partitions_zeros(r,k)
  72.   return solve(n, k)
  74. big = False
  76. funcs = Alain, Kelly, Kelly2, Dave_translated_by_Kelly
  78. if big:
  79.     funcs = funcs[1:]
  81. from functools import lru_cache, cache
  82. from itertools import accumulate
  83. from time import perf_counter as time
  84. from statistics import mean, stdev
  85. import sys
  86. import gc
  88. # 正确性
  89. for n in range(51):
  90.     for k in range(51):
  91.         expect = funcs[0](n, k)
  92.         for f in funcs[1:]:
  93.             result = f(n, k)
  94.             assert result == expect
  96. # 速度
  97. sys.setrecursionlimit(20000)
  98. times = {f: [] for f in funcs}
  99. def stats(f):
  100.     ts = [t * 1e3 for t in sorted(times[f])[:5]]
  101.     return f'{mean(ts):5.1f} ± {stdev(ts):4.1f} ms '
  102. for _ in range(25):
  103.     for f in funcs:
  104.         gc.collect()
  105.         t0 = time()
  106.         if big:
  107.            f(9000, 100)
  108.         else:
  109.             for k in range(101):
  110.                 f(100, k)
  111.         times[f].append(time() - t0)
  112. for f in sorted(funcs, key=stats):
  113.     print(stats(f), f.__name__)

请注意，由于代码段包含特定的编程注释和功能，因此某些部分的翻译可能会失去一些上下文，但我已尽力保持了代码的完整性和一致性。

英文:

Benchmark with n=100 and all k from 0 to 100, Kelly* are my solutions:

  1.   2.5 &#177;  0.1 ms  Kelly
  2. 2.8 &#177;  0.2 ms  Kelly2
  3. 3.5 &#177;  0.2 ms  Dave_translated_by_Kelly
  4. 295.0 &#177; 23.7 ms  Alain

Let c(n, k) be the number of combinations with sum n with k different numbers 1 or larger.

We get: c(n, k) = c(n-k, k) + c(n-k, k-1)

You want sum n with k different numbers 1 or larger. You either use the number 1 in the combination or you don't.

  • If you don't use 1, then you want sum n with k different numbers 2 or larger. Imagine you had such k numbers. Subtract 1 from each of them, then you have sum n-k with k different numbers 1 or larger. That's c(n-k, k).
  • If you do use 1, then you want remaining sum n-1 with k-1 different numbers 2 or larger. Imagine you had such k-1 numbers. Subtract 1 from each of them, then you have sum (n-1)-(k-1) = n-k with k-1 different numbers 1 or larger. That's c(n-k, k-1).

The faster solutions with Dave's case n=9000, k=100:

  1. 469.1 &#177;  9.2 ms  Kelly2
  2. 478.8 &#177; 17.0 ms  Kelly
  3. 673.4 &#177; 18.8 ms  Dave_translated_by_Kelly

Code (Attempt This Online!):

  1. def Kelly(n, k):
  2.     if k == 0:
  3.         return 1 if n == 0 else 0
  4.     @cache
  5.     def c(n, k):
  6.         if n &lt; k * (k+1) // 2:
  7.             return 0
  8.         if k == 1:
  9.             return 1
  10.         n -= k
  11.         return c(n, k) + c(n, k-1)
  12.     return c(n, k)
  15. # Precompute the bounds for the &quot;n &lt; ...&quot; base case
  16. def Kelly2(n, k):
  17.     if k == 0:
  18.         return 1 if n == 0 else 0
  19.     min_n_for_k = list(accumulate(range(k+1)))
  20.     @cache
  21.     def c(n, k):
  22.         if n &lt; min_n_for_k[k]:
  23.             return 0
  24.         if k == 1:
  25.             return 1
  26.         n -= k
  27.         return c(n, k) + c(n, k-1)
  28.     return c(n, k)
  31. def Alain(n, k):
  32.     @lru_cache(None)
  33.     def countNK(n,k,t=None):
  34.         t = n if t is None else t                    # track target partial sum
  35.         if k == 0:  return int(t==0)                 # empty set can sum to zero
  36.         if t &lt; 1 :  return 0                         # valid target only
  37.         if k &gt; n :  return 0                         # not enough values
  38.         return countNK(n-1,k,t)+countNK(n-1,k-1,t-n) # combine counts
  39.     return countNK(n, k)
  42. def Dave_translated_by_Kelly(n, k):
  44.   def choose(n, k):
  45.     if k &gt; n: return 0
  46.     result = 1
  47.     for d in range(1, k+1):
  48.       result *= n
  49.       result //= d
  50.       n -= 1
  51.     return result
  53.   def count_partitions_nozeroes(n, k, cache = {}):
  54.     if k==0 and n==0: return 1
  55.     if n &lt;= 0 or k &lt;= 0: return 0
  57.     # Check if the result is already memoized
  58.     if (n, k) in cache:
  59.       return cache[n, k]
  61.     # Calculate the result
  62.     result = count_partitions_nozeroes(n-1, k-1, cache) + count_partitions_nozeroes(n-k, k, cache)
  64.     # Memoize the result for future use
  65.     cache[n, k] = result
  66.     return result
  68.   def count_partitions_zeros(n,k):
  69.     return count_partitions_nozeroes(n+k, k)
  71.   def solve(n,k):
  72.     r = n - choose(k+1,2)
  73.     return count_partitions_zeros(r,k)
  75.   return solve(n, k)
  78. big = False
  80. funcs = Alain, Kelly, Kelly2, Dave_translated_by_Kelly
  82. if big:
  83.     funcs = funcs[1:]
  85. from functools import lru_cache, cache
  86. from itertools import accumulate
  87. from time import perf_counter as time
  88. from statistics import mean, stdev
  89. import sys
  90. import gc
  92. # Correctness
  93. for n in range(51):
  94.     for k in range(51):
  95.         expect = funcs[0](n, k)
  96.         for f in funcs[1:]:
  97.             result = f(n, k)
  98.             assert result == expect
  100. # Speed
  101. sys.setrecursionlimit(20000)
  102. times = {f: [] for f in funcs}
  103. def stats(f):
  104.     ts = [t * 1e3 for t in sorted(times[f])[:5]]
  105.     return f&#39;{mean(ts):5.1f} &#177; {stdev(ts):4.1f} ms &#39;
  106. for _ in range(25):
  107.     for f in funcs:
  108.         gc.collect()
  109.         t0 = time()
  110.         if big:
  111.            f(9000, 100)
  112.         else:
  113.             for k in range(101):
  114.                 f(100, k)
  115.         times[f].append(time() - t0)
  116. for f in sorted(funcs, key=stats):
  117.     print(stats(f), f.__name__)

答案4

得分: 0

让我们介绍一个函数：f(n,k,s)=从1到n中选择k个数字，它们的和为s的组合数。

为了解决这个任务，我们需要计算f(n,k,n)。

这个函数可以通过递归计算。所有的组合可以分为两组：包含最大数字和不包含最大数字的组合，这给了我们f(n,k,s)=f(n-1,k-1,s-n)+f(n-1,k,s)。递归可以在以下情况下停止：

  • n小于k -> 0（数字不够）
  • k=1，s大于n -> 0（每个数字都太小）
  • k=1，s小于1 -> 0（每个数字都太小）
  • k=1，1<=s<=n -> 1（只有一个合适的数字）
  • s小于0 -> 0

有N^2*k种可能的参数组合，所以如果我们缓存已经计算过的值，我们的复杂度将在O(N^3)范围内。

英文:

Let's introduce a function: f(n,k,s)=number of combinations of k numbers from 1 to n, having s as their sum.

To solve the task we need to calculate f(n,k,n).

The function may be calculated recursively. All combinations can be split into two groups: with and without the max number. That gives us f(n,k,s)=f(n-1,k-1,s-n)+f(n-1,k,s). The recursion may stop in the following cases:

  • n<k -> 0 (we don't have enough numbers)
  • k=1, s>n -> 0 (every number is too small)
  • k=1, s<1 -> 0 (every number is too small)
  • k=1, 1<=s<=n -> 1 (there is only one suitable number)
  • s<0 -> 0

There are N^2*k possible combinations of arguments, so if we cache the already calculated values, we will be within O(N^3).

huangapple
  • 本文由 发表于 2023年5月24日 20:53:32
  • 转载请务必保留本文链接：https://go.coder-hub.com/76323767.html
  • algorithm
  • python
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定