[HackerRank]Permuting Two Arrays

문제

 

n개의 길이를 가진 배열 q쌍과 정수 k가 주어진다. 한 쌍의 배열을 각각 A, B라고 하자. 이 배열들은 얼마든지 재정렬하거나 값들의 위치를 바꿀 수 있다. 단, 이는 한 배열 안에서만 가능하고 A, B의 값들을 서로 바꿀 수는 없다.

 

A[i]+B[i] >=k의 조건을 충족할 때 'YES', 충족하지 못할 때 'NO'를 출력한다고 할 때, 주어진 q쌍의 배열들이 각각 이 중 어디에 해당하는지를 출력하는 함수를 만들어라. 

 

단, 수의 범위 조건은 아래와 같다.

• 1≤q≤10
• 1≤n≤1000
• 1≤k≤10ⁿ(n=9)
• 0≤A[i]+B[i]≤10ⁿ(n=9)

Sample Input :

STDIN       Function
-----       --------
2           q = 2
3 10        A[] and B[] size n = 3, k = 10
2 1 3       A = [2, 1, 3]
7 8 9       B = [7, 8, 9]
4 5         A[] and B[] size n = 4, k = 5
1 2 2 1     A = [1, 2, 2, 1]
3 3 3 4     B = [3, 3, 3, 4]

Sample Output :

YES
NO

 

내 코드:

import os
import re
import sys


def twoArrays(k, A, B):
    count=0 # 기준이 될 리스트 A의 index
    i=0 # B의 index
    
    A.sort(reverse=True) # A를 내림차순으로 정렬
    B.sort() # B를 오름차순으로 정렬(가능한 작은 수로 A[count]+B[i]>=k가 충족되도록 함)
    while count<=(n-1): 
        if A[count]+B[i] >= k: # 두 값의 합이 k와 같거나 클 경우 
            B[count],B[i]=B[i],B[count] # B의 값 위치를 바꾼다
            count+=1 # 해당 값을 고정시키고 다음 index로 넘어감
            i=count

        else: # 두 값의 합이 k보다 작을 경우
            if i==(n-1):
                return 'NO'
            else: 
                i+=1 # 다음 index로 넘어감
    return 'YES' 

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    q = int(input())

    for q_itr in range(q):
        nk = input().split()

        n = int(nk[0])

        k = int(nk[1])

        A = list(map(int, input().rstrip().split()))

        B = list(map(int, input().rstrip().split()))

        result = twoArrays(k, A, B)

        fptr.write(result + '\n')

    fptr.close()

 

Comment: A[i]+B[i]>=k라는 조건을 이해하는 것이 이 문제의 핵심이라고 생각합니다. i는 배열의 인덱스를 말하는 것이고, 자연스럽게 i의 범위는 0 ≤ i ≤ n-1가 되죠. 이 말인 즉슨, A[0]+B[0], A[1]+A[1].....A[n-1]+B[n-1]이 모두 k보다 커야한다는 뜻입니다. 예시를 보면 이렇습니다. 

 

# 조건 충족('YES'가 출력) 예시
n = 3, k = 10
A = [1, 2, 3]
B = [9, 8, 7]

# A[0]+B[0]=10(k와 같거나 큼)
# A[1]+B[1]=10(k와 같거나 큼)
# A[2]+B[2]=10(k와 같거나 큼)

 

또 한 가지 중요한 점은, 'YES'와 'NO'의 충족 조건 차이를 이해하는 것입니다. 바로 위의 예시는 이렇게 나올 수도 있습니다.

 

n = 3, k = 10
A = [2, 1, 3]
B = [7, 8, 9]

# A[0]+B[0]=9(k보다 작음)
# A[1]+B[1]=9(k보다 작음)
# A[2]+B[2]=12(k와 같거나 큼)

이렇게 배열되었을 경우에는 조건을 충족하지 못하죠. 하지만 출력값은 여전히 'YES'입니다. 왜냐하면 조건이 충족되는 사례가 단 하나라도 존재하기 때문입니다. 바꿔 말하면, A[i]+B[i]>=k가 충족되는 경우가 단 하나라도 있다면 그 시점에서 출력값은 'YES'로 정해지는 것입니다. 그 앞에 아무리 충족되지 않는 경우가 많았든 간에 하등 상관이 없는 거죠. 그럼 그 하나의 'YES' 사례를 찾아낼 수 있는 함수를 만들 수 있을까요?

 

제가 생각한 과정은 이렇습니다. 우선 A든 B든 하나를 내림차순으로 정렬하고, 그 중에서 가장 큰 값(A[0] or B[0])과 더했을 때 k와 같거나 커지는 값이 있는지 찾아보는 것입니다. 

 

# 조건
n = 3, k = 10

# 주어진 배열
A = [2, 1, 3]
B = [7, 8, 9]

# B를 내림차순 정렬
B = [9, 8, 7]
A = [2, 1, 3]

 

9+2=11이므로 k보다 큽니다. 그러면 B[0], A[0]을 고정시키고 다음 값(B[1])로 넘어갑니다. B[1]+A[1]을 볼까요? 8+1=9이므로 k보다 작습니다. 그러면 A[2]로 넘어가 봅니다. B[1]+A[2], 8+3=11이므로 k보다 큽니다. 이렇게 될 경우 A[1]과 A[2]의 위치를 바꿉니다. 

 

B = [9, 8, 7]
A = [2, 3, 1]

 

어라, 그런데 충족 사례가 나오지 않는군요. B[2]+A[2], 7+1=8이므로 k보다 작습니다. 뭐가 문제일까요?

사실 9는 2가 아닌 1이었어도 조건을 충족할 수 있었습니다. 8도 마찬가지죠. 여기서 알 수 있는 점은, A의 값들이 섞여 있으면 앞에서 큰 수들끼리 더해지는 바람에 끝에서는 작은 수들끼리 짝을 짓게 되고, k보다 작아지는 상황이 생길 수 있다는 것입니다. 조건을 충족할 수 있는데도 'NO'가 출력될 수 있다는 말이죠. 

 

해결책은 간단합니다. A를 오름차순으로 정렬해주면 끝입니다.  

 

B = [9, 8, 7] # 내림차순 정렬
A = [1, 2, 3] # 오름차순 정렬

 

이렇게 정렬될 경우 가장 작은 수부터 차례로 배열을 탐색하게 되고, 위와 같은 경우가 벌어지는 것을 방지할 수 있습니다. sample input의 다른 예시로 조건에 충족되는지 확인하는 과정은 아래와 같습니다. 

 

# 다른 예시 풀이과정
n = 4, k = 5
A = [1, 2, 2, 1]
B = [3, 3, 3, 4]

# A를 내림차순 정렬, B를 오름차순 정렬
A = [2, 2, 1, 1]
B = [3, 3, 3, 4]

# A[i]과 더해 5 이상이 되는 수 탐색
A = [2, 2, 1, 1]
B = [3, 3, 4, 3]

# A[0]+B[0]>=5, 고정
A = [2, 2, 1, 1]
B = [3, 3, 3, 4]


# A[1]+B[1]>=5, 고정
A = [2, 2, 1, 1]
B = [3, 3, 3, 4]

# A[2]+B[2] < 5, A[2]+B[3]>=5, B[2]와 B[3] 위치 바꾸기(스왑)
A = [2, 2, 1, 1]
B = [3, 3, 4, 3]

# A[3]+B[3] < 5, 조건 충족 불가능



>> 'NO'