[백준(BOJ)] 3078 좋은 친구 - 스위프트 (Swift)

[Gold IV] 좋은 친구 - 3078

문제 링크

성능 요약

메모리: 75352 KB, 시간: 216 ms

분류

자료 구조, 큐, 슬라이딩 윈도우

제출 일자

2024년 1월 3일 12:34:29

문제 설명

상근이는 환갑을 바라보던 나이에 수능 시험을 다시보고 교대에 입학했고, 초등학교 선생님으로 취직했다.

• 상근: 요즘 애들은 친구를 사귀지 않나봐. 내가 앞에서 보고 있으면, 친구가 있는 학생이 별로 없는 것 같아.
• ??: 오빠! 오빠는 말콤의 친구와 성적이라는 책 안 읽어 봤어? 이 책에는 성적과 친구가 무슨 관계가 있는지 나와. 요즘 애들은 친구를 사귀기 전에 먼저 그 친구의 반 등수를 살펴봐. 말콤은 이 연구를 하기 위해서 6년동안 초등학교에서 선생님으로 위장 했었지. 하지만, 6년이라는 시간을 초등학교에서 보냈지만, 그 사람은 결국 결론을 얻지 못했어.
• 상근: 근데?
• ??: 말콤이 어느 날 자신이 초등학생이 되어 학교를 활보하는 꿈을 꾸었어. 근데 잠을 깨고 나니 내가 꿈을 꾸고 초등학생이 된건지, 아니면 초등학생이 꿈을 꾸고 지금의 내가 되어있는지를 모르겠는거야. 그래서 말콤은 상식적인 사고 방식에 큰 의문을 가졌지. 그 때 말콤은 깨달았던거야. 초등학교 친구는 부질없구나. 그제서야 알게된거야. 모든 학생은 자신과 반 등수의 차이가 K를 넘으면 친구가 아니라는거.
• 상근: 아? 근데 K는 어떻게 구해?
• ??: K는 문제에서 주어지지. 근데, 더 중요한 사실이 있어. 친구와 좋은 친구의 차이야. 말콤이 친구와 성적을 쓰고 2년 뒤에 낸 책인 좋은 친구라는 책에는 좋은 친구는 이름의 길이가 같아야 된다는 말이 나와.
• 상근: 아! 그럼 난 오늘 집에 가서 우리 반에 좋은 친구가 몇 쌍이나 있는지 구해봐야 겠어!

상근이네 반의 N명 학생들의 이름이 성적순으로 주어졌을 때, 좋은 친구가 몇 쌍이나 있는지 구하는 프로그램을 작성하시오. 좋은 친구는 등수의 차이가 K보다 작거나 같으면서 이름의 길이가 같은 친구이다.

입력

첫째 줄에 N과 K가 주어진다. (3 ≤ N ≤ 300,000, 1 ≤ K ≤ N) 다음 N개 줄에는 상근이네 반 학생의 이름이 성적순으로 주어진다. 이름은 알파벳 대문자로 이루어져 있고, 2글자 ~ 20글자이다.

출력

첫째 줄에 좋은 친구가 몇 쌍이 있는지 출력한다.

 

풀이

우선 이 문제는 제한시간이 1초이기 때문에 O(n^2)으로 풀면 시간 초과가 날 것이기 때문에 다른 방법을 생각해보고자 하였습니다. 우선 단어는 중요하지 않고 길이만 중요하기 때문에, 값들을 저장할 때에는 길이값을 저장해야겠다고 생각하였습니다.

 

예시를 통해 생각해보겠습니다.

 

 

우선 1등은 3등까지(2개 차이)가 좋은 친구이기 때문에 해당 범위만 확인해줍니다. ([3, 3, 3])

보면 2등 3등 둘 다 좋은 친구가 될 수 있기 때문에 1등은 2명의 좋은 친구쌍을 만들 수 있습니다.

 

1등을 확인했으니, 2등을 확인하기 위해 1등을 없애고 그 다음인 4등을 추가해줍니다. ([3, 3, 3])

(같은 값을 가진 배열이지만 이번 것은 2등부터 시작하여 4등까지인 배열입니다.)

위에서 확인한 바와 같이 2등 또한 2명의 좋은 친구쌍을 만들 수 있습니다.

 

3등을 확인하려고 하는데, 그 다음 등수는 없기 때문에 2등만 없애줍니다. ([3, 3])

3등은 1명의 좋은 친구쌍을 만들 수 있습니다.

 

이를 통해 n등은 해당 범위의 배열에서 본인을 제외하고 길이가 같은 원소의 개수임을 알 수 있습니다.

 

그리하여 우선 배열을 하나 생성하여 길이를 기준으로 범위 안에 있는 단어의 개수를 담아놓았습니다.

예시를 기준으로 보았을 때 wordCnt[3]의 값은 3이 되는 것입니다. (길이가 3인 단어가 3개 있기 때문)

 

여기서 처음값을 계산해줍니다. 첫 번째는 해당 배열의 원소 -1을 해줍니다.

그 다음엔 처음 값을 빼주는데 개수를 저장했기 때문에 -1을 해주면됩니다.

그리고 만약 추가될 값이 있다면 해당하는 순서에 +1(뒤에 값이 추가됨)을 해주고, 해당하는 원소 -1을 통해 쌍을 확인합니다.

(만약 해당 범위에 본인을 포함하여 2개의 길이가 같은 단어가 있다면, 1쌍이 나오는 것이기 때문에 -1을 진행하는 것입니다.)

 

이 과정을 반복하면 답을 구할 수 있습니다.

 

다음은 코드입니다.

let nk = readLine()!.split(separator: " ").map { Int(String($0))! }
let n = nk[0], k = nk[1]
var wordCnt = Array(repeating: 0, count: 21), ans = 0
var list = [Int]()

for _ in 0..<n {
    let name = readLine()!.count
    list.append(name)
}

for i in 0...k {
    wordCnt[list[i]] += 1
}
ans = wordCnt[list[0]] - 1  // 처음 값

for i in 1..<n {
    wordCnt[list[i - 1]] -= 1

    if i + k < n {  // 배열의 범위를 넘어가지 않을 떄 --> 값이 추가됨
        wordCnt[list[i + k]] += 1
    }
    ans += wordCnt[list[i]] - 1    // 쌍이니깐 -1 해주어야 함 (ex. 2개면 1쌍)
}

print(ans)