Question

我想知道是否有一种更优雅的方式来编写以下函数而不使用if语句或更少的if语句。根据我选择的参数，我希望函数返回不同的变量。

X = np.ones((3,3))
def calculate(X, output2=False, output3=False):
    y1=X*2
    y2=X*3
    y3=X*4

    if (output2==False and output3==False):
        return y1
    if (output2==True and output3==False):
        return y1,y2
    if (output2==False and output3==True):
        return y1,y3
    if (output2==True and output3==True):
        return y1,y2,y3


out1, out3=calculate(X, output3=True)

Answer 1

按元素构建返回元组元素，如果只有一个元素，则删除元组：

retval = (stuff_always_returned,)
if output2:
    retval += (more_stuff,)
if output3:
    retval += (even_more_stuff,)

if len(retval) == 1:
    return retval[0]
return retval

您也可以通过将这些计算移动到相应的if分支中来避免计算您不需要的输出部分。

retval = (stuff_always_returned,)
if output2:
    more_stuff = compute_more_stuff()
    retval += (more_stuff,)
...

Answer 2

根据掩码返回值。你可以通过迭代y-vals和mask和zip来找出要返回的内容。

def calculate(X, output2=False, output3=False):   
    ...
    mask = [True, output2, output3]
    y = [y1, y2, y3]
    return [i for i, j in zip(y, mask) if j]

out1, out3 = calculate(X, output3=True)

out1
array([[2., 2., 2.],
       [2., 2., 2.],
       [2., 2., 2.]])

out3
array([[4., 4., 4.],
       [4., 4., 4.],
       [4., 4., 4.]])

如果您只返回一个元素，则需要稍作修改。这是一个例子：

out1, *_ = calculate(X, output3=False)

Answer 3

X = np.ones((3,3))
def calculate(X, output2=False, output3=False):
    y1=X*2
    y2=X*3
    y3=X*4

    options = {(0,0):y1,
            (1,0):(y1,y2),
            (0,1):(y1,y3),
            (1,1):(y1,y2,y3)}

    return options.get((output2,output3),(None,None))

out1, out3=calculate(X, output3=True) 


[[ 2.  2.  2.]
 [ 2.  2.  2.]
 [ 2.  2.  2.]]
[[ 4.  4.  4.]
 [ 4.  4.  4.]
 [ 4.  4.  4.]]

Python中可选参数和输出的优雅编码

3 个答案: