General Question on understanding the SHIFT function in Python:
I am starting with this sample dataset:
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))
This is what I want to do in 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)
Now I am trying to do the same thing in Python.
I read that there is a SHIFT function in Python which is similar to the LEAD function in R - here is my attempt to recreate this work in 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'
})
I think this looks reasonable - but I was hoping to get a second set of eyes to verify my attempt. Any thoughts?
In 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
In 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