适用于 numpy 数组和 pandas 数据帧的列切片方法答案

【问题标题】：Column-slicing method that works on both numpy arrays and pandas dataframes适用于 numpy 数组和 pandas 数据帧的列切片方法
【发布时间】：2016-12-19 11:03:32
【问题描述】：

示例问题

假设我必须编写一个函数来返回数据表对象的第一列，但我事先不知道该对象是 numpy 二维数组还是 pandas 二维数据帧。

到目前为止尝试过

以下函数适用于 numpy 数组，但不适用于 pandas 数据帧：

def get_first_column(array_or_dataframe):
    return array_or_dataframe[:, 0]

以下函数适用于 pandas 数据帧，但不适用于 numpy 数组：

def get_first_column(array_or_dataframe):
    return array_or_dataframe.iloc[:, 0]

总结

是否可以编写一个适用于 numpy 数组和 pandas 数据帧的列切片表达式？

【问题讨论】：

检查type()并使用IF-ELSE?
为什么不直接传递df的np数组呢？例如在调用代码中：get_first_column(df.values) 基本上你要么希望类型检查逻辑进入或退出这个方法
@EdChum 如果没有与 numpy 数组和 pandas 数据帧兼容的切片表达式，这听起来确实是最简单的解决方案。
我不这么认为，这里的问题是它们是相似但不同的数据结构，所以你必须执行某种类型检查进出这个方法

标签： python arrays pandas numpy

【解决方案1】：

选项 1
np.asarray

def get_first_column(array_or_dataframe):
    return np.asarray(array_or_dataframe)[:, 0]

选项 2
try

def get_first_column(array_or_dataframe):
    try:
        return array_or_dataframe[:, 0]
    except:
        return array_or_dataframe.iloc[:, 0]

演示

df = pd.DataFrame([[1, 2], [3, 4]])

print(get_first_column(df))

print(get_first_column(df.values))

[1 3]
[1 3]

【讨论】：

聪明。谢谢！