Leetcode.315 计算右侧小于当前元素的个数【线段树】

给定一个整数数组 nums，按要求返回一个新数组 counts。数组 counts 有该性质： counts[i] 的值是  nums[i] 右侧小于 nums[i] 的元素的数量。

 

示例：

输入：nums = [5,2,6,1]
输出：[2,1,1,0] 
解释：
5 的右侧有 2 个更小的元素 (2 和 1)
2 的右侧仅有 1 个更小的元素 (1)
6 的右侧有 1 个更小的元素 (1)
1 的右侧有 0 个更小的元素
 

提示：

0 <= nums.length <= 10^5
-10^4 <= nums[i] <= 10^4

解法

1、暴力法，时间O()

2、二叉树来求和，平均时间复杂度O(nlogn), 最大时间复杂度O(),二叉树退化成链表

3、线段树，时间复杂度O(nlogn)，空间复杂度会比较大，尤其是[1,2,3,4,1212121212]这种情况

上代码

func countSmaller(nums []int) []int {
	if len(nums) == 0 {
		return nil
	}
	s := nums[0]
	e := nums[0]
	for _, n := range nums {
		if n < s {
			s = n
		}
		if n > e {
			e = n
		}
	}
	root := build(s, e)
	res := make([]int, len(nums))
	for i := len(nums) - 1; i >= 0; i-- {
		res[i] = count(root, s, nums[i]-1)
		insert(root, nums[i], 1)
	}
	return res
}

type Node struct {
	start int
	end   int
	count int
	left  *Node
	right *Node
}

func build(s, e int) *Node {
	if s > e {
		return nil
	}
	node := &Node{start: s, end: e}
	if s == e {
		return node
	}
	m := s + (e-s)/2
	node.left = build(s, m)
	node.right = build(m+1, e)
	return node
}

func insert(root *Node, n int, cnt int) {
	if root.start == n && root.end == n {
		root.count += cnt
		return
	}
	m := root.start + (root.end-root.start)/2
	var node *Node
	if n <= m {
		node = root.left
	} else {
		node = root.right
	}
	insert(node, n, cnt)
	root.count += cnt
}

func count(root *Node, s, e int) int {
	if s > e {
		return 0
	}
	if root.start == s && root.end == e {
		return root.count
	}
	m := root.start + (root.end-root.start)/2
	var c int
	if s > m {
		c = count(root.right, s, e)
	} else {
		if e > m {
			c = count(root.left, s, m) + count(root.right, m+1, e)
		} else {
			c = count(root.left, s, e)
		}
	}
	return c
}