From ddd13b616eb56ce31c8517f77969bd157898b9f2 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Tue, 14 May 2024 19:41:22 -0400
Subject: [PATCH] Add nth nearest neighbors and tests

---
 src/lib/GraphsExtensions/src/neighbors.jl |  8 ++++++++
 test/test_graphsextensions.jl             | 23 +++++++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 src/lib/GraphsExtensions/src/neighbors.jl
 create mode 100644 test/test_graphsextensions.jl

diff --git a/src/lib/GraphsExtensions/src/neighbors.jl b/src/lib/GraphsExtensions/src/neighbors.jl
new file mode 100644
index 0000000..0aa2011
--- /dev/null
+++ b/src/lib/GraphsExtensions/src/neighbors.jl
@@ -0,0 +1,8 @@
+using Graphs: AbstractGraph, neighborhood
+
+function nth_nearest_neighbors(g::AbstractGraph, v, n::Int)
+  isone(n) && return neighborhood(g, v, 1)
+  return setdiff(neighborhood(g, v, n), neighborhood(g, v, n - 1))
+end
+
+next_nearest_neighbors(g::AbstractGraph, v) = nth_nearest_neighbors(g, v, 2)
diff --git a/test/test_graphsextensions.jl b/test/test_graphsextensions.jl
new file mode 100644
index 0000000..8934e54
--- /dev/null
+++ b/test/test_graphsextensions.jl
@@ -0,0 +1,23 @@
+@eval module $(gensym())
+using NamedGraphs.NamedGraphGenerators: named_grid
+using NamedGraphs.GraphsExtensions: next_nearest_neighbors, nth_nearest_neighbors
+using Test: @test, @testset
+
+#TODO: Add tests for other graphs extensions
+@testset "GraphsExtensions" begin
+  @testset "Test nth nearest neighbours" begin
+    L = 10
+    g = named_grid((L, 1))
+    vstart = (1, 1)
+    @test only(nth_nearest_neighbors(g, vstart, L - 1)) == (L, 1)
+    @test only(next_nearest_neighbors(g, vstart)) == (3, 1)
+
+    L = 9
+    g = named_grid((L, L))
+    v_middle = (ceil(Int64, L / 2), ceil(Int64, L / 2))
+    corners = [(L, 1), (1, L), (L, L), (1, 1)]
+    @test length(next_nearest_neighbors(g, v_middle)) == 8
+    @test issetequal(nth_nearest_neighbors(g, v_middle, L - 1), corners)
+  end
+end
+end