简体   繁体   中英

R: Calculate distance between first and current row of grouped dataframe

I need to calculate the Euclidean distance between the first and current row in a dataframe. Each row is keyed by (group, month) and has a list of values. In the toy example below the key is c(month, student) and the values are in c(A, B). I want to create a distance column C, that's equal to sqrt((A_i-A_1)^2 + (B_i - B_1)^2).

So far I managed to spread my data and pull each group's first values into new columns. While I could create the formula by hand in the toy example, in my actual data I have very many columns instead of just 2. I believe I could create the squared differences within the mutate_all, and then do a row sum and take the square root of that, but no luck so far.

df <- data.frame(month=rep(1:3,2),
                 student=rep(c("Amy", "Bob"), each=3),
                 A=c(9, 6, 6, 8, 6, 9),
                 B=c(6, 2, 8, 5, 6, 7))

# Pull in each column's first values for each group
df %>% 
  group_by(student) %>% 
  mutate_all(list(first = first)) %>% 
# TODO: Calculate the distance, i.e. SQRT(sum_i[(x_i - x_1)^2]).

#Output:
  month student     A     B month_first A_first B_first
1     1 Amy         9     6           1       9       6
2     2 Amy         6     2           1       9       6
...

Desired output:

#Output:
  month student     A     B month_first A_first B_first dist_from_first
1     1 Amy         9     6           1       9       6    0
2     2 Amy         6     2           1       9       6    5
...

Here is another way using compact dplyr code. This can be used for any number of columns

df %>% 
  select(-month) %>%
  group_by(student) %>% 
  mutate_each(function(x) (first(x) - x)^2) %>%
  ungroup() %>%
  mutate(euc.dist = sqrt(rowSums(select(., -1))))

# A tibble: 6 x 4
  student     A     B euc.dist
  <chr>   <dbl> <dbl>    <dbl>
1 Amy         0     0     0   
2 Amy         9    16     5   
3 Amy         9     4     3.61
4 Bob         0     0     0   
5 Bob         4     1     2.24
6 Bob         1     4     2.24

Edit: added alternative formulation using a join. I expect that approach will be much faster for a very wide data frame with many columns to compare.

Approach 1: To get euclidean distance for a large number of columns, one way is to rearrange the data so each row shows one month, one student, and one original column (eg A or B in the OP), but then two columns representing current month value and first value. Then we can square the difference, and group across all columns to get the euclidean distance, aka root-mean-squared / RMS for each student-month.

  library(tidyverse)
  df %>% 
    group_by(student) %>% 
    mutate_all(list(first = first)) %>%
    ungroup() %>%
  # gather into long form; make col show variant, col2 show orig column
  gather(col, val, -c(student, month, month_first)) %>%
  mutate(col2 = col %>% str_remove("_first")) %>% 
  mutate(col = if_else(col %>% str_ends("_first"),
                        "first",
                        "comparison")) %>% 
  spread(col, val) %>% 
  mutate(square_dif = (comparison - first)^2) %>%
  group_by(student, month) %>%
  summarize(RMS = sqrt(sum(square_dif)))

# A tibble: 6 x 3
# Groups:   student [2]
  student month   RMS
  <fct>   <int> <dbl>
1 Amy         1  0   
2 Amy         2  5   
3 Amy         3  3.61
4 Bob         1  0   
5 Bob         2  2.24
6 Bob         3  2.24

Approach 2. Here, a long version of the data is joined to a version that is just the earliest month for each student.

library(tidyverse)
df_long <- gather(df, col, val, -c(month, student))
df_long %>% left_join(df_long %>% 
              group_by(student) %>%
              top_n(-1, wt = month) %>%
              rename(first_val = val) %>% 
              select(-month),
            by = c("student", "col")) %>%
  mutate(square_dif = (val - first_val)^2) %>%
  group_by( student, month) %>%
  summarize(RMS = sqrt(sum(square_dif)))

# A tibble: 6 x 3
# Groups:   student [2]
  student month   RMS
  <fct>   <int> <dbl>
1 Amy         1  0   
2 Amy         2  5   
3 Amy         3  3.61
4 Bob         1  0   
5 Bob         2  2.24
6 Bob         3  2.24

Instead of the mutate_all call, it'd be easier to directly calculate the dist_from_first . The only thing I'm unclear about is whether month should be included in the group_by() statement.

library(tidyverse)

df <- tibble(month=rep(1:3,2),
                 student=rep(c("Amy", "Bob"), each=3),
                 A=c(9, 6, 6, 8, 6, 9),
                 B=c(6, 2, 8, 5, 6, 7))

df%>%
  group_by(student)%>%
  mutate(dist_from_first = sqrt((A - first(A))^2 + (B - first(B))^2))%>%
  ungroup()

# A tibble: 6 x 5
#  month student     A     B dist_from_first
#  <int> <chr>   <dbl> <dbl>           <dbl>
#1     1 Amy         9     6            0   
#2     2 Amy         6     2            5   
#3     3 Amy         6     8            3.61
#4     1 Bob         8     5            0   
#5     2 Bob         6     6            2.24
#6     3 Bob         9     7            2.24

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM