LeetCode每日一题：将元素分配到两个数组中 II - 二叉索引树BIT

大家好！今天我们来聊聊一道有趣的LeetCode分配问题将元素分配到两个数组中 II。📊

问题描述

给你一个下标从 1 开始、长度为 n 的整数数组 nums 。

现定义函数 greaterCount ，使得 greaterCount(arr, val) 返回数组 arr 中 严格大于 val 的元素数量。

你需要使用 n 次操作，将 nums 的所有元素分配到两个数组 arr1 和 arr2 中。在第一次操作中，将 nums[1] 追加到 arr1 。在第二次操作中，将 nums[2] 追加到 arr2 。之后，在第 i 次操作中：

如果 greaterCount(arr1, nums[i]) > greaterCount(arr2, nums[i]) ，将 nums[i] 追加到 arr1 。
如果 greaterCount(arr1, nums[i]) < greaterCount(arr2, nums[i]) ，将 nums[i] 追加到 arr2 。
如果 greaterCount(arr1, nums[i]) == greaterCount(arr2, nums[i]) ，将 nums[i] 追加到元素数量较少的数组中。
如果仍然相等，那么将 nums[i] 追加到 arr1 。
连接数组 arr1 和 arr2 形成数组 result 。例如，如果 arr1 == [1,2,3] 且 arr2 == [4,5,6] ，那么 result = [1,2,3,4,5,6] 。

返回整数数组 result 。🙌

示例 1：

输入：nums = [2,1,3,3]
输出：[2,3,1,3]
解释：在前两次操作后，arr1 = [2] ，arr2 = [1] 。
在第 3 次操作中，两个数组中大于 3 的元素数量都是零，并且长度相等，因此，将 nums[3] 追加到 arr1 。
在第 4 次操作中，两个数组中大于 3 的元素数量都是零，但 arr2 的长度较小，因此，将 nums[4] 追加到 arr2 。
在 4 次操作后，arr1 = [2,3] ，arr2 = [1,3] 。
因此，连接形成的数组 result 是 [2,3,1,3] 。
示例 2：

输入：nums = [5,14,3,1,2]
输出：[5,3,1,2,14]
解释：在前两次操作后，arr1 = [5] ，arr2 = [14] 。
在第 3 次操作中，两个数组中大于 3 的元素数量都是一，并且长度相等，因此，将 nums[3] 追加到 arr1 。
在第 4 次操作中，arr1 中大于 1 的元素数量大于 arr2 中的数量（2 > 1），因此，将 nums[4] 追加到 arr1 。
在第 5 次操作中，arr1 中大于 2 的元素数量大于 arr2 中的数量（2 > 1），因此，将 nums[5] 追加到 arr1 。
在 5 次操作后，arr1 = [5,3,1,2] ，arr2 = [14] 。
因此，连接形成的数组 result 是 [5,3,1,2,14] 。
示例 3：

输入：nums = [3,3,3,3]
输出：[3,3,3,3]
解释：在 4 次操作后，arr1 = [3,3] ，arr2 = [3,3] 。
因此，连接形成的数组 result 是 [3,3,3,3] 。

提示：

3 <= n <= 105
1 <= nums[i] <= 109

解题思路

要解决这个问题，我们需要高效地统计数组中比当前元素大的元素数量。最直接的方式是遍历整个数组，但这样会导致时间复杂度过高，尤其是当数据量较大时。因此，我们引入了二叉索引树（Binary Indexed Tree, BIT）。🌳

为什么选择二叉索引树？

二叉索引树是一种非常适合频繁更新和查询区间和的问题的数据结构。它在以下几个方面具有优势：

1. 高效更新：每次更新操作的时间复杂度是 (O(\log n))，相比于直接遍历数组要高效得多。
2. 高效查询：每次查询操作的时间复杂度也是 (O(\log n))，能够快速获取区间和。

这些特点使得二叉索引树非常适合用于解决需要频繁统计和更新元素数量的问题。

代码实现

import scala.collection.mutable.{ArrayBuffer, Map}
import scala.util.Sorting

class BinaryIndexedTree(n: Int) {
  private val tree = new Array[Int](n + 1)

  def add(i: Int): Unit = {
    var index = i
    while (index < tree.length) {
      tree(index) += 1
      index += index & -index
    }
  }

  def get(i: Int): Int = {
    var index = i
    var sum = 0
    while (index > 0) {
      sum += tree(index)
      index -= index & -index
    }
    sum
  }
}

object Solution {
  def resultArray(nums: Array[Int]): Array[Int] = {
    val n = nums.length
    val sortedNums = nums.clone()
    Sorting.quickSort(sortedNums)

    val index = Map[Int, Int]()
    for (i <- sortedNums.indices) {
      index(sortedNums(i)) = i + 1
    }

    val arr1 = ArrayBuffer(nums(0))
    val arr2 = ArrayBuffer(nums(1))
    val tree1 = new BinaryIndexedTree(n)
    val tree2 = new BinaryIndexedTree(n)
    tree1.add(index(nums(0)))
    tree2.add(index(nums(1)))

    for (i <- 2 until n) {
      val count1 = arr1.size - tree1.get(index(nums(i)))
      val count2 = arr2.size - tree2.get(index(nums(i)))
      if (count1 > count2 || (count1 == count2 && arr1.size <= arr2.size)) {
        arr1 += nums(i)
        tree1.add(index(nums(i)))
      } else {
        arr2 += nums(i)
        tree2.add(index(nums(i)))
      }
    }

    (arr1 ++ arr2).toArray
  }
}

// 测试示例
object Main extends App {
  val nums = Array(5, 14, 3, 1, 2)
  val result = Solution.resultArray(nums)
  println(result.mkString(", "))
}

代码解析

• BinaryIndexedTree 类：实现了二叉索引树，用于高效计算比当前元素大的元素数量。
• resultArray 方法：负责将元素分配到 arr1 和 arr2 中，并返回最终结果。

时间复杂度分析 ⏱️

1. 排序：需要 (O(n \log n)) 的时间。
2. 构建映射：需要 (O(n)) 的时间。
3. 遍历和分配：每个元素插入二叉索引树和查询的操作都是 (O(\log n))，因此总时间复杂度为 (O(n \log n))。

总结

通过使用二叉索引树（BIT），我们成功地解决了这一复杂的分配问题，保证了分配的高效性和准确性。希望这个方法能帮助到大家！如果有任何问题或者建议，欢迎在评论区留言。💬

#LeetCode #算法 #数据结构 #二叉索引树 #BIT #数组分配

发布时间：2024-06-05