[python] 개미와 진딧물

문제

한 변의 길이가 N인 정사각형 모양의 땅에 개미와 진딧물이 있다. 개미는 거리 M이하에 위치한 진딧물에게서 수액을 채취할 수 있다. 이 수액이 없으면 굶주린 개미는 다른 곳으로 먹이를 찾아 떠나게 된다. 이때 거리는 맨해튼 거리이다.

 

현재 땅에 있는 개미와 진딧물의 위치가 주어졌을 때, 진딧물로부터 수액을 얻을 수 있는 개미의 수를 구해보자.

 

 

입력

첫째 줄에 정사각형의 크기 N과 거리 M이 공백을 두고 주어진다.

다음 N개의 줄에는 N개의 숫자가 공백을 두고 주어진다. 숫자는 N × N 땅을 구성하는 순서대로 주어지며, 주어지는 숫자는 0, 1, 2 중 하나이다.

0은 비어있는 위치, 1은 개미가 있는 위치, 2는 진딧물이 있는 위치임을 의미한다.

• 1 ≤ N ≤ 100
• 1 ≤ M ≤ 10
• 입력에서 주어지는 모든 수는 정수이다.

[예시]

2 0 0 0
2 0 0 1
2 0 0 0
2 0 0 1

 

 

import sys
input = sys.stdin.readline

sqrSize, distance = map(int, input().split())
locationAnt = []
locationZin = []


for i in range(sqrSize):
	line = list(map(int, input().split()))
	for j in range(sqrSize):
		if line[j] == 1:
			locationAnt.append([i, j])
		elif line[j] == 2:
			locationZin.append([i, j])

count = len(locationAnt)

for a in locationAnt:
	min_dist = 9**9
	for z in locationZin:
		dist = abs(a[0] - z[0]) + abs(a[1] - z[1])
		min_dist = min(min_dist, dist)
		
	if min_dist > distance:
		count -= 1
print(count)

 

이 문제는 조금 헷갈렸었다.

처음에는 진딧물을 기준으로 개미의 위치를 파악하여 거리가 되는 개미들을 전부 계산했었는데 테스트케이스에서 오류가 났었다.

고민해 본 결과 진딧물을 기준으로 개미 위치를 파악하면 중복되는 개미가 발생하다는 걸 깨달았다.

그래서 반대로 개미를 기준으로 진딧물의 위치를 파악하여 '전체 개미 수 - 진딧물과의 맨해튼 거리가 안 되는 개미'를 생각하며 해결했다.

 

 

 

※출처

https://multicampus-kdt.goorm.io/lecture/38996/멀티잇-코딩테스트-러닝클래스-python-6월반