How to make my function take less time in python?

Question

I am trying to do this problem:

https://www.chegg.com/homework-help/questions-and-answers/blocks-pyramid-def-pyramidblocks-n-m-h-pyramid-structure-although-ancient-mesoamerican-fam-q38542637

This is my code:

def pyramid_blocks(n, m, h):
    return sum((n+i)*(m+i) for i in range(h))

But the problem is that whenever I try to test it, it tells me that it is too slow, so is there a way to make it take less time? I also tried to do it with lists, but with that too it takes too much time.

def pyramid_blocks(n, m, h):
    return sum((n+i)*(m+i) for i in range(h))

def pyramid_blocks(n, m, h):
    r=[]
    t=[]
    mlp=[]
    for i in range(h):
        r.append(n+i)
        t.append(m+i)
    for i, j in zip(r,t):
        mlp.append(i*j)
    return sum(mlp)

Asked By: sfreg

||

Source

Answer 1

You should use math to understand the components that are part of the answer. you need to calculate sum of all numbers until h, denote s, and the sum of all squares 1^2+2^2+3+2+... denotes sum_power
it comes down to: h*m*n + s*n + s*m + sum_power

this is a working solution:

import time


def pyramid_blocks(n, m, h):
    return sum((n+i)*(m+i) for i in range(h))


def efficient_pyramid_blocks(n, m, h):
    s = (h-1)*h//2
    sum_power = (h-1)*h*(2*(h-1)+1)//6
    return h*n*m + s*(n+m) + sum_power


if __name__ == '__main__':
    n, m, h = 2123377, 2026271, 2437
    start = time.time()
    print(pyramid_blocks(n, m, h))
    print(time.time()-start)
    start = time.time()
    print(efficient_pyramid_blocks(n, m, h))
    print(time.time() - start)

and the output is:

333723479800000
0.016993045806884766
333723479800000
1.2159347534179688e-05

Answered By: Ohad Zadok

Answer 2

A bit of math can help:
here we use a symbolic calculator to get the equation for the blocks in the pyramid using a sigma sum of (n+count)*(m+count) from count = 0 to count = h-1

Then we just paste the equation into code (the double slash // is just to use integer division instead of float division):

def count_blocks(n,m,h):
    return (2*(h**3)+3*(h**2)*m+3*(h**2)*n+6*h*m*n-3*(h**2)-3*h*m-3*h*n+h)//6

Edit: This also outputs the correct result of 10497605327499753 for n,m,h = (2123377, 2026271, 2437)

Answered By: Some nerd who does not have a

How to make my function take less time in python?

Question:

Answers: