Python使用numpy排序后去除array()标识返回列表元组的方法

问题根源 #

• xsort和ysort本身是numpy数组类型，直接返回会保留numpy的array()标识
• 原有返回逻辑用了[xsort]和[ysort]，额外给数组套了一层列表嵌套，和预期的单层列表结构不符
• 额外兼容问题：开头的解包写法([a,b,c],[d,e,f]) = data仅支持长度为3的输入列表，不具备通用性

修复方案 #

最简修复（仅适配现有固定长度逻辑） #

将原函数末尾的return语句修改为：

return (xsort.tolist(), ysort.tolist())

.tolist()方法可以直接把一维numpy数组转为普通Python列表，同时去掉外层多余的列表包裹，直接返回两个列表组成的元组即可匹配预期输出。

通用优化版本（支持任意长度的输入列表） #

如果需要适配任意长度的输入列表，可调整完整代码如下：

import numpy as np

def sort_data(data):
    """ (tuple) -> tuple
    
    data is a tuple of two lists.
    
    Returns a copy of the input tuple sorted in
    non-decreasing order with respect to the 
    data[0]
    
    >>> sort_data(([5, 1, 7], [1, 2, 3]))
    ([1, 5, 7], [2, 1, 3])

    >>> sort_data(([2, 4, 8], [1, 2, 3]))
    ([2, 4, 8], [1, 2, 3])

    >>> sort_data( ([11, 4, -5], [1, 2, 3]))
    ([-5, 4, 11], [3, 2, 1])
    """
    x, y = data
    xarray = np.array(x)
    yarray = np.array(y)
    sort_idx = np.argsort(xarray)
    xsort = xarray[sort_idx].tolist()
    ysort = yarray[sort_idx].tolist()
    return (xsort, ysort)

验证结果 #

调用测试用例sort_data(([5, 1, 7], [1, 2, 3]))，输出结果为([1, 5, 7], [2, 1, 3])，完全符合预期。

内容的提问来源于stack exchange，提问作者ashnotallyson