Faster version of raster::union for a single SpatialPolygonsDataFrame

self_union.R

self_union <- function(x) {
  stopifnot(requireNamespace("rgeos"))

  n <- length(x)
  if (n < 2) {
    return(x)
  }

  if (.hasSlot(x, 'data')) {
    x <- as(x, 'SpatialPolygons')
  }

  x_poly_ids <- vapply(slot(x, "polygons"), function(y) slot(y, "ID"), character(1))

  x <- SpatialPolygonsDataFrame(x, data.frame(dummy = 1:n, row.names = x_poly_ids))

  ints <- get_intersections(x)
  if (length(ints) == 0) return(x)

  # n_ints <- length(ints)
  int_ids <- unique(c(names(ints), unlist(sapply(ints, names, USE.NAMES = FALSE))))
  no_int_ids <- setdiff(x_poly_ids, int_ids)

  x_ints <- x[int_ids, ]
  x_no_ints <- x[no_int_ids,]

  no_ints_df <- make_identity_df(no_int_ids)

  x_no_ints <- SpatialPolygonsDataFrame(x_no_ints, no_ints_df)

  u <- x_ints[int_ids[1], ]
  names(u) <- paste0("ID.", int_ids[1])
  for (id in int_ids[2:length(int_ids)]) {
    z <- x_ints[id, ]
    names(z) <- paste0("ID.", id)
    u <- union(u, z)
  }

  all_list <- c(u, x_no_ints, recursive = TRUE)

  all_list <- rebuild_poly_ids(all_list)

  all_list <- align_colnames(all_list)

  unioned_all <- do.call("rbind", all_list)

  unioned_all@data$dummy <- NULL
  unioned_all@data[is.na(unioned_all@data)] <- 0
  unioned_all@data[unioned_all@data > 0] <- 1

  unioned_all$count <- rowSums(unioned_all@data)

  unioned_all
}

make_identity_df <- function(names) {
  mat <- diag(nrow = length(names))
  colnames(mat) <- paste0('ID.', names)
  data.frame(mat, row.names = names)
}

get_intersections <- function(x) {
  int_matrix <- rgeos::gIntersects(x, byid = TRUE, returnDense = TRUE)

  int_matrix[lower.tri(int_matrix, diag = TRUE)] <- NA

  int_list <- apply(int_matrix, 1, which)

  int_list[lapply(int_list, length) > 0]

}

rebuild_poly_ids <- function(x) {
  lens <- vapply(x, length, integer(1))
  i <- 1
  for (l in seq_along(x)) {
    this_end <- i + lens[l] - 1
    x[[l]] <- spChFIDs(x[[l]], as.character(i:this_end))
    i <- this_end + 1
  }
  x
}

align_colnames <- function(x) {
  names_list <- lapply(x, function(y) colnames(y@data))
  all_names <- Reduce(base::union, names_list)
  ret <- lapply(x, function(y) {
    missing_cols <- setdiff(all_names, colnames(y@data))
    y@data <- cbind(y@data, make_empty_df(missing_cols, nrow(y@data)))
    y
  })
  ret
}

make_empty_df <- function(colnames, nrows) {
  ncol <- length(colnames)
  df <- data.frame(matrix(NA_integer_, ncol = ncol, nrow = nrows))
  setNames(df, colnames)
}

test.R

library(raster)
library(rgdal)
library(rgeos)

p1 <- Polygon(cbind(c(2,4,4,1,2),c(2,3,5,4,2)))
p2 <- Polygon(cbind(c(5,5,3,3,5),c(2,4,3,2,2)))
p3 <- Polygon(cbind(c(7,7,6,6,7), c(2,4,3,2,2)))
ps1 <- Polygons(list(p1), "s1")
ps2 <- Polygons(list(p2), "s2")
ps3 <- Polygons(list(p3), "s3")
spp <- SpatialPolygons(list(ps1, ps2, ps3), 1:3); plot(spp)

spdf <- SpatialPolygonsDataFrame(spp, data.frame(A = c(1,2,3), B = letters[8:10],
                                                 row.names = c("s1", "s2", "s3")))

self_union(spdf)