[Python] 分段函数

遇上了一个很简单的小问题，就是写一个分段函数。分段函数如图1：

图1. 分段函数

一开始写了个简单版本log_norm0，只能逐元素一个一个得进行。不用想，对元素个数多的向量，肯定慢成乌龟。
后来想到使用一个指示函数，来区分分段的两种情况，就得到log_norm1。不过这种情况比较特殊，不是每次都能成功构造的。
最后，查到numpy中有函数piecewise(x, condlist, funclist, *args, **kw)，它是专门用来构造分段函数，x是输入，condlist表示分段的条件，funclist就表示对应分段的处理函数。这就得到了log_norm2。

# elementwise
def log_norm0(x):
    if x >= 0:
        return np.log(x + 1)
    else:
        return - np.log(- x + 1)

# indicator
def log_norm1(x):
    # ind = np.where(x > 0, 1, 0)
    ind = (x > 0)
    return np.log(x * ind + 1) - np.log(- x * (1.0 - ind) + 1)

# numpy.piecewise()
def log_norm2(x):
    return np.piecewise(x, [x >= 0, x < 0], [lambda x: np.log(1 + x), lambda x: - np.log(1 - x)])

最后，观察一个各个函数的运行时间。

tic = time.time()
for i in range(x.size):
    y[i] = log_norm0(x[i])
toc = time.time()
print('log0: ', toc - tic)

tic = time.time()
y = log_norm1(x)
toc = time.time()
print('log1: ', toc - tic)

tic = time.time()
z = log_norm2(x)
toc = time.time()
print('log2: ', toc - tic)

观察结果，还是使用指示函数的方法最快，不过跟piecewise差别不大。

log0:  33.59282732009888
log1:  0.4863457679748535
log2:  0.5942573547363281

参考资料：

1. https://blog.csdn.net/shu15121856/article/details/76080060
2. https://docs.scipy.org/doc/numpy/reference/generated/numpy.piecewise.html

扩展资料（待填坑）：
https://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.call.html
https://docs.scipy.org/doc/numpy/reference/generated/numpy.frompyfunc.html
https://blog.csdn.net/kezunhai/article/details/46127845