fork download
  1. from collections import defaultdict
  2.  
  3. dp = defaultdict(lambda: -1)
  4. prefixsum = []
  5. trace = []
  6.  
  7. def getSubArraySum(i, j):
  8. if i == 0:
  9. return prefixsum[j]
  10. return (prefixsum[j] - prefixsum[i - 1])
  11.  
  12. def rec(cur, arr, k):
  13. if cur >= len(arr):
  14. return 0
  15. if dp[cur] != -1:
  16. return dp[cur]
  17.  
  18. s1 = -1; s2 = -1
  19. # If we choose the subarray starting at `cur`
  20. if cur + k - 1 < len(arr):
  21. s1 = getSubArraySum(cur, cur + k - 1) + rec(cur + k + 1, arr, k)
  22. # If we ignore the subarray starting at `cur`
  23. s2 = rec(cur + 1, arr, k)
  24.  
  25. dp[cur] = max(s1, s2)
  26.  
  27. if s1 >= s2:
  28. trace[cur] = (True, cur + k + 1)
  29. return s1
  30. trace[cur] = (False, cur + 1)
  31. return s2
  32.  
  33. def getTrace(arr, trace, k):
  34. itr = 0
  35. subArrays = []
  36. while itr < len(trace):
  37. if trace[itr][0]:
  38. subArrays.append(arr[itr : itr + k])
  39. itr = trace[itr][1]
  40. return subArrays
  41.  
  42. def solve(arr, k):
  43. global dp, trace, prefixsum
  44.  
  45. dp = defaultdict(lambda: -1)
  46. trace = [(False, 0)] * len(arr)
  47. prefixsum = [0] * len(arr)
  48.  
  49. prefixsum[0] = arr[0]
  50. for i in range(1, len(arr)):
  51. prefixsum[i] += prefixsum[i - 1] + arr[i]
  52.  
  53. print("Array :", arr)
  54. print("Max sum: ", rec(0, arr, k))
  55. print("Subarrays: ", getTrace(arr, trace, k))
  56. print("-- * --")
  57.  
  58. solve([1, 2, 3, 4], 3)
  59. solve([1, 2, 3, 1, 7, 9] , 2)
  60. solve([1, 5, 3, 1, 7, 9] , 1)
Success #stdin #stdout 0.02s 27664KB
stdin
Standard input is empty
stdout
Array : [1, 2, 3, 4]
Max sum:  9
Subarrays:  [[2, 3, 4]]
-- * --
Array : [1, 2, 3, 1, 7, 9]
Max sum:  21
Subarrays:  [[2, 3], [7, 9]]
-- * --
Array : [1, 5, 3, 1, 7, 9]
Max sum:  15
Subarrays:  [[5], [1], [9]]
-- * --