fork download
  1. class Solution {
  2. private var tree: [Int] = []
  3.  
  4. // Building the segment tree
  5. private func buildTree(_ start: Int, _ end: Int, _ p: Int) -> Int {
  6. if start == end {
  7. tree[p] = 1
  8. return tree[p]
  9. }
  10. let mid = (start + end) / 2
  11. tree[p] = buildTree(start, mid, 2 * p + 1) + buildTree(mid + 1, end, 2 * p + 2)
  12. return tree[p]
  13. }
  14.  
  15. // Comparator function for sorting
  16. private static func cmp(_ a: [Int], _ b: [Int]) -> Bool {
  17. if a[0] == b[0] {
  18. return a[1] > b[1]
  19. }
  20. return a[0] < b[0]
  21. }
  22.  
  23. // Function to get the index using the segment tree
  24. private func getIndex(_ start: Int, _ end: Int, _ value: inout Int, _ p: Int, _ index: inout Int) {
  25. if start > end { return }
  26. if start == end {
  27. tree[p] -= 1
  28. index = start
  29. return
  30. }
  31.  
  32. let mid = (start + end) / 2
  33. if tree[2 * p + 1] > value {
  34. tree[p] -= 1
  35. getIndex(start, mid, &value, 2 * p + 1, &index)
  36. } else {
  37. tree[p] -= 1
  38. value -= tree[2 * p + 1]
  39. getIndex(mid + 1, end, &value, 2 * p + 2, &index)
  40. }
  41. }
  42.  
  43. // Reconstruct the queue using the segment tree
  44. func reconstructQueue(_ people: [[Int]]) -> [[Int]] {
  45. var people = people
  46. people.sort(by: Solution.cmp)
  47.  
  48. let n = people.count
  49. var size = 1
  50. while size < n {
  51. size *= 2
  52. }
  53. let arrSize = size - 1
  54. size *= 2
  55. tree = Array(repeating: 0, count: size)
  56.  
  57. buildTree(0, n - 1, 0)
  58. var res = Array(repeating: [Int](), count: n)
  59.  
  60. for person in people {
  61. var index = 0
  62. var k = person[1]
  63. getIndex(0, n - 1, &k, 0, &index)
  64. res[index].append(person[0])
  65. res[index].append(person[1])
  66. }
  67.  
  68. return res
  69. }
  70. }
Success #stdin #stdout 0.01s 6160KB
stdin
// Example usage
let solution = Solution()
let people = [[7, 0], [4, 4], [7, 1], [5, 0], [6, 1], [5, 2]]
let result = solution.reconstructQueue(people)
print(result)
stdout
Standard output is empty