NumPy sum、mean、cumsum 及参数 axis，keepdims，where 的含义

许多数理统计函数都可以在 NumPy 中使用，既可以针对数组中所有数据进行计算，也可以沿某个维度进行计算。

其中一类函数我们称为 aggregations，也经常叫作 reductions，例如 sum，mean，std，这些函数的返回结果要比原始 array 的维度少。

我们有如下 array：

import numpy as np

arr = np.random.randn(5, 4)
arr
"""
array([[ 0.0099062 ,  0.7066875 ,  0.11536878, -1.56797423],
       [-0.08768619, -0.36325377,  0.9233594 , -0.93038662],
       [ 0.57972329, -1.08741835,  0.93362563, -1.45548137],
       [ 0.16445882, -0.15527661,  0.20953435, -0.10089461],
       [-0.39885311, -1.56403318, -0.29668125,  0.02043629]])
"""
1
2
3
4
5
6
7
8
9
10
11

对数组中所有数据求均值：两种写法，直接调用 array 方法或者使用函数

arr.mean()
"""
-0.21724195193767884
"""
1
2
3
4

np.mean(arr)
"""
-0.21724195193767884
"""
1
2
3
4

对数组中所有数据求和：

arr.sum()
"""
-4.344839038753577
"""
1
2
3
4

我们也可以通过 axis 参数指定这类函数沿哪个维度进行计算。例如，我们将 mean 中的 axis 指定为 1，那么将会沿 axis 1 方向计算均值，即计算每行的均值：

arr.mean(axis=1)
"""
array([-0.18400294, -0.1144918 , -0.2573877 ,  0.02945549, -0.55978281])
"""
1
2
3
4

将 sum 中的 axis 指定为 0，那么将会沿 axis 0 方向求和，即对每列求和：

arr.sum(axis=0)
"""
array([ 0.26754901, -2.46329442,  1.88520691, -4.03430054])
"""
1
2
3
4

axis 的方向示意如下：

sum、mean 还有一个参数 keepdims，如果设置为 True，那么被 reduced 的轴（也就是我们设置的 axis 值）就会被保留，大小变为 1：

arr.sum(axis=0, keepdims=True)
"""
array([[ 0.26754901, -2.46329441,  1.88520691, -4.03430054]])
"""

arr.sum(axis=0, keepdims=True).shape
"""
(1, 4)
"""
1
2
3
4
5
6
7
8
9

我们还可以通过参数 where 来指定我们想要包括在内计算的元素。where 是一个布尔类型的 array：

arr.sum(axis=0, where=[True, False, False, True]) # 只计算第一列和最后一列
"""
array([ 0.26754901,  0.        ,  0.        , -4.03430054])
"""
1
2
3
4

a = np.array([[5, 9, 13], [14, 10, 12], [11, 15, 19]])
a.mean()
"""
12.0
"""

a.mean(where=[[True], [False], [False]]) # 只计算第一行的均值
"""
9.0
"""
1
2
3
4
5
6
7
8
9
10

这里的 where 参数使用了广播机制。例如我们在 sum 中传入的 [True, False, False, True] 会被广播为：

array([[ True, False, False,  True],
       [ True, False, False,  True],
       [ True, False, False,  True],
       [ True, False, False,  True],
       [ True, False, False,  True]])
1
2
3
4
5

与 arr 的形状匹配。

---

还有一类函数，它们的返回结果和原始 array 有相同的大小。如计算累加值的函数 cumsum：

arr = np.array([0, 1, 2, 3, 4, 5, 6])
arr.cumsum()
"""
array([ 0,  1,  3,  6, 10, 15, 21])
"""
1
2
3
4
5

对于多维 array，也可以指定 axis，表示在特定维度进行运算：

arr = np.arange(9).reshape((3, 3))
arr
"""
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
"""
1
2
3
4
5
6
7

axis = 0 表示计算每列各自的累加值，但中间结果会保留，不会 aggregate 成一个值。因此和原始 array 有相同大小：

arr.cumsum(axis=0)
"""
array([[ 0,  1,  2],
       [ 3,  5,  7],
       [ 9, 12, 15]])
"""
1
2
3
4
5
6

---

References

Python for Data Analysis, 2${}^{\mathrm{n}\mathrm{d}}$ edition. Wes McKinney.