leetcode day 29: 01 matrix

Transform an arbitrarily-shaped rectangular matrix of zeroes and ones into a map of distance from the nearest zero, using a NSEW neighborhood.

My first impulse was, “oh, find all the zeroes, and then work out from each one to fix the ones”, followed by “that sounds hard to do, I’ll try the recurrence relation and caching”. The recurrence relation and caching simply does not work because the problem space is too big. Even with caching the intermediate results I still ran out of memory because the neighborhood processing recursed 4x for every cell.

So I went back to the original thought. I didn’t get much of anywhere until I got the hint that what I needed to do was mark all the 1’s as 999999 (an impossible value, marking them as “not done”). After that, I pulled up my memories of linked list processing and it was actually kind of easy! We start up by creating a QueueItem struct:

type QueueItem struct {
    i int,
    j int,
    next *QueueItem
}

Now we need to find all the zeros and add them to the queue, and mark all the ones as “not done yet”. In hindsight, I should have made the queue operations at least a separate sub — mostly because there was a test case with a single zero, which messed up my enqueue/dequeue logic! — but they’re inline:

    m := len(mat)
    n := len(mat[0])

    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            if mat[i][j] == 0 {
                item := &queueItem{
                    i: i,
                    j: j,
                    next: queue,
                }
                if queue == nil {
                    // tail stays here while we add to the head
                    tail = item
                }
                queue = item
            } else {
                mat[i][j] = 9999999
            }
        }
    }

Pretty straightforward — now we have all the zeroes on the queue, and can work our way out from each one. The 999999 is a value guaranteed greater than the largest possible real value in the map. Now all we have to do is process the neighborhoods around each zero; if we find a cell we need to update, we know from the definition of the problem that these are one away from the cell we took off the queue, so for the zero cells, these will all be ones. We set the value in the matrix, and then we add the cell we just processed to the tail of the queue. This way we process all the zeroes, then all the ones, and so on until we’ve processed every cell that needs it. When that finishes, we’re done and can return the updated matrix.

    neighborhood := [][]int{ {1,0}, {-1,0}, {0, -1}, {0, 1}}

    for queue != nil {
        // next cell in priority popped off the quque
        cell := queue
        queue = queue.next
        for _, dir := range neighborhood {
            // Figure out the next cell in the neighborhood of the one we
            // pulled off the queue.
            offsetI, offsetJ := cell.i + dir[0], cell.j + dir[1]

            // If we're in bounds...
            if offsetI >= 0 && offsetI < m && offsetJ >= 0 && offsetJ < n {
                // ...and the cell we're looking at has a value > the current
                // cell + 1, update it to the value of the current cell + 1
                // and queue it. Note that neighbors closer to a zero or 
                // zeroes themselves will be fail the test and will be
                // ignored; they're already done.
                if mat[offsetI][offsetJ] > mat[cell.i][cell.j] + 1 {
                    newCell := &queueItem{
                        i: offsetI,
                        j: offsetJ,
                    }
                    // This is the special case that bit me: it's possible
                    // that the queue emptied. Reset the head and tail 
                    // pointers to the cell we're adding.
                    if queue == nil {
                        queue = newCell
                        tail = newCell
                    } else {
                        // Didn't empty, so add to the tail. This guarantees
                        // that we process all items closer to a zero first.
                        tail.next = newCell 
                        tail = tail.next
                    }
                    // This updates the cell to the right distance and
                    // removes the "unprocessed" marker.                    
                    mat[offsetI][offsetJ] = mat[cell.i][cell.j] + 1
                }
            }
        }
    }

    // when we arrive here there are no more queue items to process,
    // we've marked all the cells and are done.
    return mat

So that’s it! The priority queue is stupid easy to do after the other list exercises, and I should have trusted my initial idea and tried it first. The problem I got at Bloomberg was actually a variation on this, and I could solve it fairly fast now.

I will have to figure out dynamic programming at some point, but this problem didn’t need it.