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

时间:2016-12-19 11:03:32

标签: python arrays pandas numpy

示例问题

假设我必须编写一个返回数据表对象第一列的函数,但我事先并不知道这个对象是一个numpy二维数组还是一个pandas二维数据帧。 / p>

到目前为止已经尝试了

以下函数适用于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数据帧上的列切片编写单个表达式?

1 个答案:

答案 0 :(得分: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]