[codeforces] 1408B - Array Sum

문제 링크: http://codeforces.com/contest/1408/problem/B

 

음이 아닌 정수로 이루어진, 감소하지 않는 순서로 정렬된 리스트 a_1, a_2, ..., a_n이 주어져있을 때, 크기 n짜리 리스트 m개를 만들어서 각 리스트의 i번째 아이템을 다 더해서 a_i가 되도록 만들고자 한다. 이때, m개의 리스트들 각각은, 최대 k개의 서로 다른 원소로 이루어질 수 있다.

 

예를 들어, n = 5, k = 3, a = [1, 2, 3, 4, 5]로 주어져있다고 하자. a는 [1, 2, 0, 0, 0], [0, 0, 3, 4, 0], [0, 0, 0, 0, 5]의 3개의 리스트들을 같은 위치에 있는 원소들끼리 더해서 만들 수 있다. [1, 2, 0, 0, 0], [0, 0, 3, 4, 5]는, 두 번째 리스트가 서로 다른 4개의 원소를 들고 있으므로 k = 3 조건을 만족하지 않는다. 이때, a는 [0, 1, 2, 2, 2], [1, 1, 1, 2, 3]의 두 리스트를 더해서 만드는 것이 가능하고, 더 적은 개수의 리스트를 더해서 만드는 것은 불가능하므로, m값의 최소값은 2임을 알 수 있다.

 

n, k, a값이 주어질때, 최소 m의 값을 출력하는 것이 문제.

 

a에 들어있는 서로 다른 원소 개수가 x개일때, 서로 다른 k개의 원소로 이루어진 리스트를 뺄 때마다 서로 다른 원소 개수가 최대 k-1개 줄어들 수 있다는 아이디어로 문제에 접근했다. 예를 들어, a=[1, 2, 3, 4, 5], k=3으로 주어졌을 때, [0, 0, 0, 1, 2]를 빼면 [1, 2, 3, 3, 3]으로, 뒤에 있는 4와 5를 3으로 맞출 수 있다. 이런 시행은 서로 같은 원소의 개수가 몇 개가 있더라도 영향을 받지 않음을 알 수 있는데, [1, 2, 3, 3, 3, 4, 4, 5, 5]가 주어져도 [0, 0, 0, 0, 0, 1, 1, 2, 2]를 빼면 4, 5가 3으로 변하고, 시행 회수는 그대로 1회이다.

 

다시 [1, 2, 3, 4, 5] 예시로 돌아와보자. 처음에는 서로 다른 원소 개수가 5개이고, [0, 0, 0, 1, 2]를 빼서 서로 다른 원소 개수를 3개로 줄였다. 그러면, 비슷한 방법으로 시행을 한 번 더 하면 서로 다른 원소의 개수가 3->1개로 줄어들 것이다. 그러면 서로 다른 원소 개수가 1개인, 즉 모든 아이템이 같은 값으로 이루어진 리스트가 남을 것이고, 여기에 앞서 두 번의 시행에 쓴 2개의 리스트를 더하면 원래 리스트가 나오므로, 총 3개의 리스트를 더하면 된다- 하는 결론을 내릴 수도 있지만, 잘 생각해보면 3->1개로 서로 다른 원소 개수를 줄일 필요 없이 남은 리스트 자체가 '서로 다른 3개의 원소로 이루어져있다'는 조건을 만족하므로, 3->1 시행을 할 필요 없이 계산이 끝남을 알 수 있다. 즉, 답을 구하고 싶으면 a를 이루는 서로 다른 원소의 개수가 x개일때, ceil((x-1)/(k-1))값을 구하면 된다.

 

k=1일때 나누기 오류가 날 것이므로 k=1인 케이스를 생각해서 예외처리 해주면 된다고 생각해서 아래와 같이 답안을 제출했다.

from math import ceil
 
for _ in range(int(input())):
    n,k = map(int,input().split())
    a=list(map(int,input().split()))
    x = len(set(a))
    
    if k==1:
        if x==1: print(1)
        else: print(-1)
    else:
        print(ceil((x-1)/(k-1)))

그리고 시스 채점에서 혼났다...

 

서로 다른 원소 개수 x값이 1이 되는 케이스에서 해당 계산 결과는 어떻게 될까? x=1이므로 ceil((x-1)/(k-1))값을 계산하면 0이 되는데, 이건 0개의 리스트를 가지고 a 리스트를 만들 수 있다는 뜻으로, 이것도 예외처리를 해주었어야 하는 것. x==1인 케이스를 따로 elif문으로 빼서 print(1)을 해주어도 되고, 마지막 else문에 max(1, 계산식) 하는 식으로, 결과가 0이 되는 케이스가 나와도 1을 출력하도록 해주어도 될 것이다. 다시 제출해서 맞았다..