문제 내용은 간단하다. 나와 K 순위내에 있는 이름글자가 동일한 친구는 몇명일까? 이걸 모든 친구들에 대해서 구해주면 된다. (그냥 구하면 N^2 시간복잡도로 시간 초과가 난다.)
기본 아이디어
- 이름 글자수가 20자로 정해져 있으니 크기 20의 배열 생성
- 위에서 선언한 배열 중 이름 글자수에 해당하는 곳에 순위를 저장(나는 리스트형태로 추가함) -> 동일한 길이의 이름을 가진 학생들이 모이게됨
- 저장된 순위를 오름차순으로 정렬(이미 정렬된 형태로 들어오는데 왜 정렬하는지 모르겠음…)
- 이진탐색을 통해 선택된 학생과 K순위 내에 있는 학생수 구한다. (이진탐색시 동일한 값에 대해 앞에 넣을지 뒤에 넣을지 부등호 주의해야 한다.)
내소스가 최선은 아니며 누군가에게 도움이 되길 바라는 마음으로 작성한다.
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.StringTokenizer;
public class Main {
static int[] data;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
int N = Integer.parseInt(st.nextToken());
int K = Integer.parseInt(st.nextToken());
List<Integer>[] arr = new ArrayList[21];
for(int i=0;i<21;i++) {
arr[i] = new ArrayList<Integer>();
}
for(int i=0;i<N;i++) {
String name = br.readLine();
arr[name.length()].add(i);
}
for(int i=0;i<21;i++) {
Collections.sort(arr[i]);
}
long ans=0;
for(int i=0;i<21;i++) {
data = new int[arr[i].size()+1];
for(int j=0;j<arr[i].size();j++) {
data[j] = arr[i].get(j);
}
data[arr[i].size()] = Integer.MAX_VALUE;
for(int j=0;j<data.length-1;j++) {
int idx = binarySearch(j, data.length, data[j]+K);
if(j==data.length-2) {
continue;
}
if(idx > j) {
ans += idx-j-1;
}
}
}
System.out.println(ans);
}
private static int binarySearch(int stIdx, int edIdx, int val) {
int mid =0, ret = 0;
while(stIdx <= edIdx) {
mid = (stIdx+edIdx)/2;
if(data[mid]<=val) { //부등호 주의(같은 값일때 뒤에 넣음)
stIdx = mid+1;
}else {
edIdx = mid-1;
ret = mid;
}
}
return ret;
}
}