假设我必须编写一个返回数据表对象第一列的函数,但我事先并不知道这个对象是一个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数据帧上的列切片编写单个表达式?
答案 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]