Find next element in matrix - swift

Question

I have simple two dimensional array with different number of elements in each row. My goal is find next element with elegant solution.
If we have the last element we should get the first one.
[
[1, 0, 234, 345, 678],
[123, 456, 789],
[0, 9]
]
678 -> 123
9 -> 1

Answer 1

If mat is a nested array then mat.joined() is a “lazy flat collection” of all the nested elements. That suggest to solve the problem for arbitrary collections:
Given a collection and an element of that collection, return the immediate successor of the element. The collection is viewed as cyclic, so that the successor of the last element is the first element.
A straight-forward implementation is (explanations inline):
extension Collection where Element: Comparable {
func cyclicNext(after elem: Element) -> Element? {
// Find (first) index of given element
guard var idx = firstIndex(of: elem) else {
return nil // Element not present in collection
}
formIndex(after: &idx) // Advance index by one
if idx == endIndex { // If past the end ...
idx = startIndex // ... then jump back to the start
}
return self[idx]
}
}
If the element occurs more than once in the collection then the successor of the first occurrence is returned.
This can be applied to your matrix:
let mat: [[Int]] = [[1, 0, 234, 345, 678],
[123, 456, 789],
[0, 9]]
print(mat.joined().cyclicNext(after: 1)) // 0
print(mat.joined().cyclicNext(after: 678)) // 123
print(mat.joined().cyclicNext(after: 123)) // 456
print(mat.joined().cyclicNext(after: 9)) // 1
print(mat.joined().cyclicNext(after: 999)) // nil
It can also be applied to other collections:
print([1, 3, 5, 7].cyclicNext(after: 3)) // 5
print([1, 3, 5, 7].cyclicNext(after: 7)) // 1
print("Hello World".cyclicNext(after: "W")) // "o"
print("Hello World".cyclicNext(after: "d")) // "H"

Answer 2

I believe this does do the job :
extension Array where Element == [Int] {
func element(after x: Int) -> Int? {
var arrayIndex = 0
while arrayIndex < self.count {
//If an array contains the searched element
if let xIndex = self[arrayIndex].firstIndex(where: { $0 == x }) {
//if the next element is in the same array
if xIndex < self[arrayIndex].count - 1 {
return self[arrayIndex][xIndex + 1]
}
//if the next element is in the next array
else if arrayIndex < self.count - 1 {
return self[arrayIndex + 1][0]
}
//if the x is the last element in the last array
else {
return self[0][0]
}
}
arrayIndex += 1
}
return nil
}
}
And here are some test cases :
let mat = [
[1, 0, 234, 345, 678],
[123, 456, 789],
[0, 9]
]
mat.element(after: 678) //123
mat.element(after: 9) //1
mat.element(after: 1) //0
mat.element(after: 0) //234
mat.element(after: 3) //nil

Answer 3

A possible solution is to flatten your array.
Tested on Playground:
let mat: [[Int]] = [[1, 0, 234, 345, 678],
[123, 456, 789],
[0, 9]]
func next(element: Int, in matrix: [[Int]]) -> Int? {
let flatten = matrix.flatMap({ $0 })
guard let index = flatten.firstIndex(of: element) else {
print("Element \(element) not found in matrix")
return nil
}
if index == flatten.endIndex - 1 {
return flatten.first
} else {
return flatten[index + 1 ]
}
}
let result1 = next(element: 678, in: mat)
print("End of subarray Test: \(result1)")
let result2 = next(element: 9, in: mat)
print("Last one Test: \(result2)")
let result3 = next(element: 66, in: mat)
print("Not present test: \(result3)")
let result0 = next(element: 1, in: mat)
print("First one test: \(result0)")
Output:
$>End of subarray Test: Optional(123)
$>Last one Test: Optional(1)
$>Element 66 not found in matrix
$>Not present test: nil
$>First one test: Optional(0)
I don't know if it's enough elegant for you.
One optimization would have to keep the flatten and not recalculate each time.