[Java+Python]2805번, K번째 수

목차

이진 탐색

 

[Java+Python]1920번, 수 찾기

[Java+Python]1654번, 랜선 자르기

[Java+Python]2805번, 나무 자르기

[Java+Python]2110번, 공유기 설치

[Java+Python]12015번, 가장 긴 증가하는 부분 수열 2

문제

 

세준이는 크기가 N×N인 배열 A를 만들었다. 

배열에 들어있는 수 A[i][j] = i×j 이다. 

이 수를 일차원 배열 B에 넣으면 B의 크기는 N×N이 된다. 

B를 오름차순 정렬했을 때, B[k]를 구해보자.

배열 A와 B의 인덱스는 1부터 시작한다.

 

입력

 

첫째 줄에 배열의 크기 N이 주어진다. 

N은 $10^5$보다 작거나 같은 자연수이다. 

둘째 줄에 k가 주어진다. 

k는 min($10^9, N^2$) 보다 작거나 같은 자연수이다.

 

출력

 

B[k]를 출력한다.

 

풀이

 

주어진 크기(N)의 이차원 배열에서 K번째 수를 찾는 문제이다.

 

주의할 점은 배열의 인덱스가 1로 시작한다는 것과, 해당 배열을 B에 넣은 뒤 오름차순으로 정렬했다는 점이다.

 

바로 이 부분 덕분에 이 문제를 이진 탐색으로 풀 수 있게 된다.

 

또한 배열에 들어가는 숫자도 n이 주어진 순간 결정되기 때문에 따로 선언할 필요가 없다.

 

문제 풀이의 알고리즘은 다음과 같다.

 

1. n, k를 입력받는다.
2. left = 1, right = n x n, answer = 0으로 초기화한다.
3. 이진 탐색을 시작한다. 이진 탐색은 아래와 같이 진행한다.
   
   
   1. 중간값 mid = (left + right) / 2로, count = 0으로 초기화한다.
   2. 이어서 1부터 n까지 반복문을 돌며 (mid / i)와 n 중 작은 값으로 count를 갱신한다.
      여기서 mid / i가 등장한 이유는 위와 같은 특별한 형태의 배열에서 mid / i가 mid 이하의 개수를 나타내기 때문이다.
   3. count >= k 라면 answer = mid로, right = mid - 1로 갱신하고,
      그렇지 않다면 left = mid + 1로 갱신한 뒤 이진 탐색의 처음으로 돌아간다.
4. 위와 같이 이진탐색이 끝난 뒤에 answer에 저장된 값을 출력한다.

여기서 mid / i가 등장하는 이유를 조금 더 자세히 짚고 넘어가자면,

 

우리가 가지고 있는 배열은 모든 행이 i의 배수로 이루어져 있다는 것을 확인할 수 있다.

 

즉, 1번째 행은 1의 배수, 2번째 행은 2의 배수... 와 같은 식이다.

 

이어서 우리는 mid가 k번째 값이 될 수 있는지 확인해야 하며, 이를 위해 mid보다 작거나 같은 값의 개수를 구해야 한다.

 

여기서 mid / i의 등장 이유가 발생하는데, 위에 적었듯이 각 행은 i의 배수로 이루어져 있기 때문에

 

mid를 i로 나누면 mid보다 작거나 같은 i의 배수가 몇 개 존재하는지 알 수 있게 되는 것이다.

 

추가로 해당 값을 n과 비교하는 이유는 모든 행에서의 mid / i의 값은 n을 초과할 수 없기 때문이며,

 

이처럼 각 행마다 mid 이하의 값이 몇 개나 있는지 세어보면 전체 배열에서 mid 이하의 값을 알 수 있게 된다.

 

이어서 이 정보를 근거로 이진 탐색을 진행하는 것이다.

 

Java

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

public class Prob1300 {

	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		long n = Integer.parseInt(br.readLine());
		long k = Integer.parseInt(br.readLine());

		long left = 1;
		long right = n * n;
		long answer = 0;

		while (left <= right) {
			long mid = (left + right) / 2;
			long count = 0;

			for (int i = 1; i <= n; i++) {
				count += Math.min(mid / i, n);
			}

			if (count >= k) {
				answer = mid;
				right = mid - 1;
			} else {
				left = mid + 1;
			}
		}

		System.out.println(answer);
	}
}

 

Python

 

import sys

n = int(sys.stdin.readline().rstrip())
k = int(sys.stdin.readline().rstrip())

left = 1
right = n**2
answer = 0

while left <= right:
    mid = (left + right) // 2
    count = 0

    for i in range(1, n + 1):
        count += min(mid // i, n)

    if count >= k:
        answer = mid
        right = mid - 1
    else:
        left = mid + 1

print(answer)

 

Performance