[백준] 1009번 분산처리 - 파이썬

https://www.acmicpc.net/problem/1009

문제

재용이는 최신 컴퓨터 10대를 가지고 있다. 어느 날 재용이는 많은 데이터를 처리해야 될 일이 생겨서 각 컴퓨터에 1번부터 10번까지의 번호를 부여하고, 10대의 컴퓨터가 다음과 같은 방법으로 데이터들을 처리하기로 하였다.

1번 데이터는 1번 컴퓨터, 2번 데이터는 2번 컴퓨터, 3번 데이터는 3번 컴퓨터, ... ,

10번 데이터는 10번 컴퓨터, 11번 데이터는 1번 컴퓨터, 12번 데이터는 2번 컴퓨터, ...

총 데이터의 개수는 항상 ab개의 형태로 주어진다. 재용이는 문득 마지막 데이터가 처리될 컴퓨터의 번호가 궁금해졌다. 이를 수행해주는 프로그램을 작성하라.

입력

입력의 첫 줄에는 테스트 케이스의 개수 T가 주어진다. 그 다음 줄부터 각각의 테스트 케이스에 대해 정수 a와 b가 주어진다. (1 ≤ a < 100, 1 ≤ b < 1,000,000)

출력

각 테스트 케이스에 대해 마지막 데이터가 처리되는 컴퓨터의 번호를 출력한다.

예제 입력

5
1 6
3 7
6 2
7 100
9 635

예제 출력

1
7
6
1
9

처음 시도한 풀이

def sol(base, expo):
    current = (base * base) % 10
    endNum = [base, current]
    count = 1
    while current != base:
        current = (current * base) % 10
        if current not in endNum:
            endNum.append(current)
        count += 1

    return endNum[(expo % count) - 1]


for i in range(int(input())):
    a, b = map(int, input().split())
    print(sol(a, b))

count: a를 몇 번 제곱해야 일의 자리 수가 a의 일의 자리 수와 같아지는 지를 나타내는 횟수

endNum: a를 계속 제곱하면서 반복되는 일의 자리 수 리스트

endNum 리스트에서 (b를 count로 나눈 나머지 - 1) 번째 수가 정답

이렇게 했지만 시간초과가 났다.

정답 풀이

endNumList = [[10], [1], [2, 4, 8, 6], [3, 9, 7, 1], [4, 6],
              [5], [6], [7, 9, 3, 1], [8, 4, 2, 6], [9, 1]]


def sol(base, expo):
    baseEndNum = base % 10
    expoEndNum = expo % len(endNumList[baseEndNum])
    return endNumList[baseEndNum][expoEndNum - 1]


for i in range(int(input())):
    a, b = map(int, input().split())
    print(sol(a, b))

endNum 리스트를 매 케이스마다 계산하지 않고 처음부터 0~9까지에 대해서 endNum 리스트를 미리 만들어 두면 여러 번 계산하지 않아도 된다.