[코드트리] 악수와 전염병의 상관관계 2

https://www.codetree.ai/missions/5/problems/correlation-between-shaking-hands-and-infectious-diseases2?&utm_source=clipboard&utm_medium=text

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

< 입력 형식 >

첫 번째 줄에 정수 N, K, P, T가 각각 공백을 사이에 두고 주어집니다.

두 번째 줄부터는 T개의 줄에 걸쳐 각 줄마다 t, x, y에 대한 정보가 공백을 사이에 두고 주어집니다. 이는 t초에 x 개발자와 y 개발자가 악수를 나눴음을 의미하고, x, y 값은 항상 다르게 주어짐을 가정해도 좋습니다. 또한, 입력으로 주어지는 t값은 모두 다름을 가정해도 좋습니다.

• 2 ≤ N ≤ 100
• 1 ≤ K ≤ 250
• 1 ≤ P ≤ N
• 1 ≤ T ≤ 250
• 1 ≤ t ≤ 250

< 풀이 >

입력을 받을 준비를 한다.

n, k, p, T = map(int, input().split())

 

입력 예제를 보니 악수를 나눈 정보가 시간순으로 들어온다는 보장이 없다는 것을 체크한다.

4 2 2 3
7 1 2 
5 2 3
6 2 4

 

이 경우 정확한 추적을 위해 시간 순으로 악수를 나눈 정보를 정리할 필요가 있다.

별도의 배열을 선언해서 t에 대한 오름차순으로 정렬하자.

logs = []
for _ in range(T):
    t, x, y = map(int, input().split())
    logs.append((t, x, y))

logs.sort(key=lambda x: x[0])

 

t를 보니 최댓값으로 250을 가진다. 따라서 개발자마다 timestamp을 갖고 있다면 어떨까 생각해보았다.

즉 각각의 개발자는 251초의 시간에 따른 정보(양성인지 음성인지)를 기억할 수 있는 timestamp를 갖고 있는 것이다.

이차원 배열로 선언해준다.

arr = [[0 for _ in range(251)] for _ in range(n+1)]

 

p번째 개발자가 최초 감염자이므로 0초에 p번째 개발자의 timestamp에 양성 반응을 찍는다.

#최초 감염자
arr[p][0] = 1

 

k번 악수를 하면 더 이상 전염력이 없다는 것에 유의한다.

멀쩡한 개발자가 최초 감염되는 경우는 감염된 개발자 입장에서 악수한 횟수로 카운트 하지 않는다.

아래 예시를 보자.

#n k p T
8 2 1 5
1 1 2 
2 2 1
3 2 4
4 2 5
5 5 1

 

입력 예시가 위와 같이 주어졌다고 가정하자.

초기 (t=0) 배열 상황은 이렇다.

개발자 1이 최초 감염자이다.

[
  1  [1, 0, 0, 0, 0, 0, 0, 0, ...],
  2  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  3  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  4  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  5  [0, 0, 0, 0, 0, 0, 0, 0, ...],
    ...
]

 

t=1일 때, 개발자 1과 개발자 2가 악수를 했다

[
  1  [1, 1, 0, 0, 0, 0, 0, 0, ...],
  2  [0, 1, 0, 0, 0, 0, 0, 0, ...],
  3  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  4  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  5  [0, 0, 0, 0, 0, 0, 0, 0, ...],
    ...
]

 

t=2일 때, 개발자 2와 개발자 1이 악수를 했다.

[
  1  [1, 1, 1, 0, 0, 0, 0, 0, ...],
  2  [0, 1, 1, 0, 0, 0, 0, 0, ...],
  3  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  4  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  5  [0, 0, 0, 0, 0, 0, 0, 0, ...],
    ...
]

 

t=3일 때, 개발자 2와 개발자 4가 악수를 했다.

개발자 2의 배열은 1이 3개 이지만 맨 앞의 1은 감염됐을 때 생긴 것이므로 개발자 2가 악수를 한 횟수에 포함되지 않는 것이다.

따라서 개발자 4에게도 전염이 된다

[
  1  [1, 1, 1, 0, 0, 0, 0, 0, ...],
  2  [0, 1, 1, 1, 0, 0, 0, 0, ...],
  3  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  4  [0, 0, 0, 1, 0, 0, 0, 0, ...],
  5  [0, 0, 0, 0, 0, 0, 0, 0, ...],
    ...
]

 

t=4일 때, 개발자 2와 개발자 5가 악수를 했다.

개발자 2는 악수 횟수(k=2)를 소진했기 때문에 이제 전염성이 없다.

개발자 5를 전염시키지 못한다

[
  1  [1, 1, 1, 0, 0, 0, 0, 0, ...],
  2  [0, 1, 1, 1, 0, 0, 0, 0, ...],
  3  [0, 0, 0, 0, 0, 0, 0, 0, ...],
  4  [0, 0, 0, 1, 0, 0, 0, 0, ...],
  5  [0, 0, 0, 0, 0, 0, 0, 0, ...],
    ...
]

 

정리하면 각 개발자 배열이 갖고 있는 1의 개수가 k+1개가 되면 전염성을 상실한다

코드로 옮겨보자

for log in logs:
    t, x, y = log
    
    # 둘중 하나가 감염자이면서 악수를 k번보다 적게 했다면 감염
    if (0 < arr[x].count(1) < k+1 or 0 < arr[y].count(1) < k+1):
        arr[x][t] = 1
        arr[y][t] = 1

 

감염자 여부는 timestamp를 돌아가면서 1이 있는지만 확인하면 된다.

# 1이 하나라도 있으면 감염된 사람임
result = []
for developer in arr:
    if developer.count(1) > 0:
        result.append(1)
    else:
        result.append(0)

 

결과를 출력한다

for i in range(1, len(result)):
    print(result[i], end="")

[ 전체 코드 ]

n, k, p, T = map(int, input().split())

arr = [[0 for _ in range(251)] for _ in range(n+1)]

#최초 감염자
arr[p][0] = 1

logs = []
for _ in range(T):
    t, x, y = map(int, input().split())
    logs.append((t, x, y))

logs.sort(key=lambda x: x[0])

for log in logs:
    t, x, y = log

    # 둘중 하나가 감염자이면서 악수를 k번보다 적게 했다면 감염
    if (0 < arr[x].count(1) < k+1 or 0 < arr[y].count(1) < k+1):
        arr[x][t] = 1
        arr[y][t] = 1

# 1이 하나라도 있으면 감염된 사람임
result = []
for developer in arr:
    if developer.count(1) > 0:
        result.append(1)
    else:
        result.append(0)

for i in range(1, len(result)):
    print(result[i], end="")

< 느낀 점 >

전체에 대해서 K번의 전염 횟수가 제한이 있는게 아니라 전염병에 걸린 개발자 각각에 대해서 K번의 전염 횟수에 제한이 있다는 것을 일찍 알았다면 삽질을 덜 했을것 같다.