shift() 函数 (Python) 和 lead() 函数 (R) 之间的等价性

Question

关于理解 Python 中的 SHIFT 函数的一般问题：

我从这个示例数据集开始：

start = structure(list(id = c(111L, 111L, 111L, 111L, 222L, 222L, 222L
), year = c(2010L, 2011L, 2012L, 2013L, 2018L, 2019L, 2020L),
col2 = c("A", "BB", "A", "C", "D", "EE", "F"), col3 = c(242L,
213L, 233L, 455L, 11L, 444L, 123L), col4 = c(1213L, 5959L,
9988L, 4242L, 333L, 1232L, 98L)), class = "data.frame", row.names = c(NA,
-7L))

这就是我想在 R 中做的事情：

library(dplyr)

end <- start %>%
mutate(year_end = lead(year),
     col2_end = lead(col2),
     col3_end = lead(col3),
     col4_end = lead(col4)) %>%
mutate_at(vars(ends_with("_end")), ~ifelse(is.na(.), "END", .)) %>%
rename(year_start = year,
     col2_start = col2,
     col3_start = col3,
     col4_start = col4)

现在我正尝试在 Python 中做同样的事情。

我读到 Python 中有一个 SHIFT 函数，它类似于 R 中的 LEAD 函数 - 这是我在 Python 中重新创建这项工作的尝试：

import pandas as pd

start = pd.DataFrame({
    'id': [111, 111, 111, 111, 222, 222, 222],
    'year': [2010, 2011, 2012, 2013, 2018, 2019, 2020],
    'col2': ['A', 'BB', 'A', 'C', 'D', 'EE', 'F'],
    'col3': [242, 213, 233, 455, 11, 444, 123],
    'col4': [1213, 5959, 9988, 4242, 333, 1232, 98]
})

end = start.assign(
    year_end=start['year'].shift(-1),
    col2_end=start['col2'].shift(-1),
    col3_end=start['col3'].shift(-1),
    col4_end=start['col4'].shift(-1)
).fillna('END')

end = end.rename(columns={
    'year': 'year_start',
    'col2': 'col2_start',
    'col3': 'col3_start',
    'col4': 'col4_start'
})

我认为这看起来很合理——但我希望得到另一双眼睛来验证我的尝试。有什么想法吗？

Answer 1

在R中：

start %>%
   mutate(across(-id, ~lead(as.character(.x),default = 'END'),.names = '{col}_end'))

   id year col2 col3 col4 year_end col2_end col3_end col4_end
1 111 2010    A  242 1213     2011       BB      213     5959
2 111 2011   BB  213 5959     2012        A      233     9988
3 111 2012    A  233 9988     2013        C      455     4242
4 111 2013    C  455 4242     2018        D       11      333
5 222 2018    D   11  333     2019       EE      444     1232
6 222 2019   EE  444 1232     2020        F      123       98
7 222 2020    F  123   98      END      END      END      END

在 python 中：

(start.join(start.drop('id', axis = 1).shift(-1, fill_value = 'END')
    .add_suffix('_end')))
 
    id  year col2  col3  col4 year_end col2_end col3_end col4_end
0  111  2010    A   242  1213     2011       BB      213     5959
1  111  2011   BB   213  5959     2012        A      233     9988
2  111  2012    A   233  9988     2013        C      455     4242
3  111  2013    C   455  4242     2018        D       11      333
4  222  2018    D    11   333     2019       EE      444     1232
5  222  2019   EE   444  1232     2020        F      123       98
6  222  2020    F   123    98      END      END      END      END