SW Expert Academy 5648. [모의 SW 역량테스트] 원자 소멸 시뮬레이션

문제 링크: https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AWXRFInKex8DFAUo&categoryId=AWXRFInKex8DFAUo&categoryType=CODE

SW Expert Academy

처음에 0.5초를 처리하기 위해 2000*2000 크기의 배열을 선언하였다가 시간 초과가 났다. 다른 사람들의 코드를 리뷰해서 파이썬에 있는 자료구조 중 하나인 딕셔너리를 사용하여 문제를 해결하였다. 해결한 코드는 아래와 갔다.

t = int(input())
for test in range(1, t + 1):
    n = int(input())
    atoms = []
    for i in range(n):
        x, y, d, e = map(int, input().split())
        x *= 2
        y *= 2
        atoms.append([x, y, d, e])
    dx = [0, 0, -1, 1]
    dy = [1, -1, 0, 0]
    count = 0
    energy = 0
    while 1:
        dic = {}
        for a in atoms:
            a[0] += dx[a[2]]
            a[1] += dy[a[2]]
            try:
                dic[(a[0], a[1])].append(a)
            except:
                dic[(a[0], a[1])] = [a]
        atoms = []
        for i in dic.keys():
            if len(dic[i]) > 1:
                for atom in dic[i]:
                    energy += atom[3]
                    count += 1
            else:
                x = dic[i][0][0]
                y = dic[i][0][1]
                if -4001 < x < 4001 and -4001 < y < 4001:
                    atoms.append(dic[i][0])
                else:
                    count += 1
        if count == n:
            break
    print('#{} {}'.format(test, energy))