如何组合两个几乎相同的功能但有效地返回不同的 output

Question

I have two functions, which are almost same, in a main function. For a simplified example,

def func1(a):
    for i in range(999):
       a += np.exp(-i)     # some massive calculation
    return a

def func2(a):
    b = 0
    for i in range(999):
       a += np.exp(-i)     # same massive calculation
       b += np.exp(-(a+b)) # additional massive calculation 
    return a, b

def main():
    a = 1
    for _ in range(9999):
        a = func1(a)
    print(a)
    a = 0
    c = 0
    for _ in range(9999):
        a, b = func2(a)
        c +=b
    print(a, c)

I know I could have combined func1 and func2 as

def func3(a, return_b):
    b = 0
    for i in range(999):
       a += np.exp(-i) 
       if return_b:        # add if statement
           b += np.exp(-(a+b))            
    return a,b  

and the main function becomes

def main():
    a = 1
    for _ in range(9999):
        a,_ = func3(a, False)
    print(a)
    a = 0
    c = 0
    for _ in range(9999):
        a, b = func3(a, True)
        c +=b
    print(a, c)

But I think it's really expensive to add if statement in a large number of loops. And I would rather repeat func1 in func2 and leave both functions in the code. I wonder what the best way is to combine func1 and func2 together by considering the running time and code re-use efficiency.

Answer 1

This code below is less expensive than if inside for loop and as you wanted it still is in one function.

def func3(a, return_b):
    b = 0
    #checks only once
    if return_b:
        for i in range(999):
            a += np.exp(-i) 
            b += np.exp(-(a+b))
    else:
        for i in range(999):
            a += np.exp(-i) 
    return a,b