Search code examples
algorithmmathdynamic-programmingprefix-sum

How to solve M times prefix sum with better time complexity


The problem is to find the prefix sum of array of length N by repeating the process M times. e.g.

Example N=3
M=4
array = 1 2 3
output = 1 6 21
Explanation:
Step 1 prefix Sum = 1 3 6
Step 2 prefix sum = 1 4 10
Step 3 prefix sum = 1 5 15
Step 4(M) prefix sum = 1 6 21

Example 2:
N=5
M=3
array = 1 2 3 4 5
output = 1 5 15 35 70

I was not able to solve the problem and kept getting lime limit exceeded. I used dynamic programming to solve it in O(NM) time. I looked around and found the following general mathematical solution but I still not able to solve it because my math isn't that great to understand it. Can someone solve it in a better time complexity?

https://math.stackexchange.com/questions/234304/sum-of-the-sum-of-the-sum-of-the-first-n-natural-numbers


Solution

  • Hint: 3, 4, 5 and 6, 10, 15 are sections of diagonals on Pascal's Triangle.

    JavaScript code:

    function f(n, m) {
      const result = [1];
      
      for (let i = 1; i < n; i++)
        result.push(result[i-1] * (m + i + 1) / i);
      
      return result;
    }
    
    console.log(JSON.stringify(f(3, 4)));
    console.log(JSON.stringify(f(5, 3)));