문제풀이
- 1 ~ N값에 대해 그룹을 지어서 idxList를 만든다.
- 각 그룹의 idxList의 누적합 리스트를 만든다.
- 누적합을 이용하여 개수를 카운팅하고, 같은 그룹의 3개의 숫자가 묶인 경우를 뺀다.
주의사항
- overflow 조심
소요시간
15분
package AtCoder.ProblemList.Difficulty800_1199.Sandwiches
fun main() {
val N = readln().trim().toInt()
val A = readln().trim().split(" ").map(String::toInt)
val idxLists = List(N + 1) { mutableListOf<Int>() }
A.forEachIndexed { index, i -> idxLists[i].add(index) }
val idxSumLists = List(N + 1) { LongArray(idxLists[it].size) }
for (i in idxLists.indices) {
for (j in idxLists[i].indices) {
idxSumLists[i][j] = idxSumLists[i].getOrElse(j - 1) { 0 } + idxLists[i][j]
}
}
var result = 0L
for (i in idxLists.indices) {
for (j in idxLists[i].indices) {
result += (idxLists[i][j] - 1).toLong() * j - idxSumLists[i].getOrElse(j - 1) { 0 }
if (j >= 2) result -= j.toLong() * (j - 1) / 2
}
}
println(result)
}