[Python] 2차원 리스트 원소의 상하좌우 탐색

Python의 알고리즘 문제들을 풀다보면 2차원 리스트의 원소의 상하좌우 등을 탐색해야하는 경우가 있다. 이때, list indexing 이용하면 쉽게 탐색할 수 있다.

>>> n = int(input())
>>> lists = [list(map(int, input().split())) for _ in range(n)]

Python의 2차원 리스트를 위와 같이 만들었을때 2행 2열의 원소인 39에 접근할려면 전체 리스트에서 2번째 행을 indexing한 후 두번째 원소를 indexing하면 된다.

>>> print(lists[1][1])
39

따라서 다음과 같은 논리를 통해 지정 원소의 상하좌우를 탐색할 수 있다.

원소 - 39의 index: (1, 1)
상 - 13의 index: (0, 1)
하 - 25의 index: (2, 1)
좌 - 12의 index: (1, 0)
우 - 30의 index: (1, 2)

이것을 함수화 하면

원소index: (x, y)
상 index: (x-1, y)
하 index: (x+1, y)
좌 index: (x, y-1)
우 index: (x, y+1)

예를 들면 list[x][y] 의 위의 원소는 list[x-1][y]를 사용하면 확인할 수 있다.

Python 상하좌우를 순서대로 list로 만들어 반환하는 함수

def find(lists, index:set):
    x, y = index
    dx = [-1, 1, 0, 0]
    dy = [0, 0, -1, 1]
    result = []
    for i in range(4):
        a = lists[x+dx[i]][y+dy[i]]
        result.append(a)
    return result

만든 함수를 이용해 위의 2차원 리스트에 적용해보기

>>>result = find(lists, (1,1))
>>>print(result)
[13, 25, 12, 30]

주의:

위함수는 2차원 리스트의 테두리는 고려하고 만들지 않은것이기 때문에 find함수에 테두리 원소의index를 사용하면 잘못된 값을 반환한다.