[CodeSignal] 배수를 만드는 순서쌍 찾기

배열 a와 정수 k를 받을 때 i < j 인 i, j 에 대하여 a[i] + a[j] 가 정수 k 의 배수가 되는 i, j를 찾는 문제이다. 

어렵지 않은 구현 문제라고 생각하여 이러한 코드를 짜고 제출했다.

 

def solution(a, k):
    cnt = 0
    for i in range(len(a)-1):
        for j in range(i+1, len(a)):
            if (a[i] + a[j])%k == 0:
                cnt += 1
    return cnt

 

아주 간단하게 생각할 수 있는 완전탐색 코드이다. 주어진 배열에서 나올 수 있는 모든 경우의 수를 탐색하는 코드인데, 보기 좋게 시간초과 오류가 났다. 제한 시간이 4초이기에 여유로울줄 알았는데 입력 크기가 엄청난가보다.

 

아무리 생각해도 발상이 안떠올라서 전지전능하신 ChatGPT의 도움을 받아 다음의 코드를 짰다.

 

def solution(a, k):
    
    reminders = [0 for i in range(k)]
    
    for i in range(len(a)):
        reminders[int(a[i]%k)] += 1

    cnt = (reminders[0] * (reminders[0]-1)) // 2

    for i in range(1, (k-1) // 2 + 1):
        cnt += reminders[i] * reminders[k - i]
        
    if k%2 == 0:
        cnt += (reminders[k//2]*(reminders[k//2]-1))//2
    
    return cnt

 

시간복잡도에 대해 최적화 과정이 들어간 코드이다. 코드는 크게 4부분으로 나뉜다.

 

(1) k 길이의 배열을 만들고 각 a의 숫자를 k로 나누고 나머지의 빈도를 계산한다.

(2) (나머지가 0인 경우) 중복되지 않게 숫자 두개를 선택하는 조합

(3) (나머지가 0이 아닌 경우) 나머지가 i인 경우에서 하나, k-i 인 경우에서 하나 선택하는 조합

(4) (k가 짝수인 경우) 나머지가 k // 2 인 숫자를 겹치지 않게 두개 선택하는 조합

 

이를 모두 더하면 구하고자 하는 cnt 가 나온다.

 

이런 방법으로 시간복잡도를 O(n^2)에서 O(n)까지 줄일 수 있다.