|
| 1 | +package g0401_0500.s0480_sliding_window_median |
| 2 | + |
| 3 | +// #Hard #Array #Hash_Table #Heap_Priority_Queue #Sliding_Window |
| 4 | +// #2023_01_01_Time_409_ms_(100.00%)_Space_44.6_MB_(81.48%) |
| 5 | + |
| 6 | +import java.util.TreeSet |
| 7 | + |
| 8 | +class Solution { |
| 9 | + fun medianSlidingWindow(nums: IntArray, k: Int): DoubleArray { |
| 10 | + require(k >= 1) { "Input is invalid" } |
| 11 | + val len = nums.size |
| 12 | + val result = DoubleArray(len - k + 1) |
| 13 | + if (k == 1) { |
| 14 | + for (i in 0 until len) { |
| 15 | + result[i] = nums[i].toDouble() |
| 16 | + } |
| 17 | + return result |
| 18 | + } |
| 19 | + val comparator = Comparator { a: Int?, b: Int? -> |
| 20 | + if (nums[a!!] != nums[b!!] |
| 21 | + ) Integer.compare(nums[a], nums[b]) else Integer.compare(a, b) |
| 22 | + } |
| 23 | + val smallNums = TreeSet(comparator.reversed()) |
| 24 | + val largeNums = TreeSet(comparator) |
| 25 | + for (i in 0 until len) { |
| 26 | + if (i >= k) { |
| 27 | + removeElement(smallNums, largeNums, i - k) |
| 28 | + } |
| 29 | + addElement(smallNums, largeNums, i) |
| 30 | + if (i >= k - 1) { |
| 31 | + result[i - (k - 1)] = getMedian(smallNums, largeNums, nums) |
| 32 | + } |
| 33 | + } |
| 34 | + return result |
| 35 | + } |
| 36 | + |
| 37 | + private fun addElement(smallNums: TreeSet<Int?>, largeNums: TreeSet<Int?>, idx: Int) { |
| 38 | + smallNums.add(idx) |
| 39 | + largeNums.add(smallNums.pollFirst()!!) |
| 40 | + if (smallNums.size < largeNums.size) { |
| 41 | + smallNums.add(largeNums.pollFirst()) |
| 42 | + } |
| 43 | + } |
| 44 | + |
| 45 | + private fun removeElement(smallNums: TreeSet<Int?>, largeNums: TreeSet<Int?>, idx: Int) { |
| 46 | + if (largeNums.contains(idx)) { |
| 47 | + largeNums.remove(idx) |
| 48 | + if (smallNums.size == largeNums.size + 2) { |
| 49 | + largeNums.add(smallNums.pollFirst()!!) |
| 50 | + } |
| 51 | + } else { |
| 52 | + smallNums.remove(idx) |
| 53 | + if (smallNums.size < largeNums.size) { |
| 54 | + smallNums.add(largeNums.pollFirst()) |
| 55 | + } |
| 56 | + } |
| 57 | + } |
| 58 | + |
| 59 | + private fun getMedian(smallNums: TreeSet<Int?>, largeNums: TreeSet<Int?>, nums: IntArray): Double { |
| 60 | + return if (smallNums.size == largeNums.size) { |
| 61 | + (nums[smallNums.first()!!].toDouble() + nums[largeNums.first()!!]) / 2 |
| 62 | + } else nums[smallNums.first()!!].toDouble() |
| 63 | + } |
| 64 | +} |
0 commit comments