CODEWITHSUNDEEP
python
Benjy Kessler
Benjy Kessler

Reputation: 7616

Finding valid neighbor indices in 2d array

I have a 2d list in Python. Given an index I want to find all the neighbors of that index. So if my list is 3x3 then given the index (1, 1) I want to return [(0, 1), (1, 0), (2, 1), (1, 2), (2, 2), (0, 0), (0, 2), (2, 0)] but if the index is (0, 0) then I want to only return [(0, 1), (1,0), (1, 1)]. I know how to do this with ugly if statements. My question is, is there a pretty Pythonic magic one liner for this?

Upvotes: 2

Views: 518

Answers (2)

B. M.
B. M.

Reputation: 18628

a constant time solution for a 3x3 space for example, with list comprehensions :

valid={(x,y) for x in range(3) for y in range (3)} 
dirs=[(dx,dy) for dx in (-1,0,1) for dy in (-1,0,1) if (dx,dy)!=(0,0)] 
def voisins(x,y): return [(x+dx,y+dy) for (dx,dy) in dirs  if (x+dx,y+dy) in valid]

Upvotes: 3

Marius
Marius

Reputation: 60060

How ugly is your current code? I can't think of any obvious ways to do this automagically, just doing it manually without trying to write fancy one liners I don't think it looks too bad:

def find_neighours(index, x_size, y_size):
    neighbours = []
    x1, y1, = index
    for x_move in (-1, 0, 1):
        for y_move in (-1, 0, 1):
            if x_move == 0 and y_move == 0:
                continue
            x2, y2 = x1 + x_move, y1 + y_move
            if x2 < 0 or x2 >= x_size:
                continue
            if y2 < 0 or y2 >= y_size:
                continue
            neighbours.append((x2, y2))
    return neighbours

Outputs:

find_neighours((1, 1), 3, 3)
Out[2]: [(0, 0), (0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1), (2, 2)]

find_neighours((0, 0), 3, 3)
Out[3]: [(0, 1), (1, 0), (1, 1)]

Upvotes: 1

Related Questions