315. Count of Smaller Numbers After Self #

Problem #

You are given an integer array nums and you have to return a new counts array. The counts array has the property where counts[i] is the number of smaller elements to the right of nums[i].

Example:

Input: [5,2,6,1]
Output: [2,1,1,0] 
Explanation:
To the right of 5 there are 2 smaller elements (2 and 1).
To the right of 2 there is only 1 smaller element (1).
To the right of 6 there is 1 smaller element (1).
To the right of 1 there is 0 smaller element.

Problem Summary #

Given an integer array nums，return a new array counts as required。The array counts has this property： the value of counts[i] is  the number of elements to the right of nums[i] that are smaller than nums[i]。

Example:

Input: [5,2,6,1]
Output: [2,1,1,0] 
Explanation:
To the right of 5 there are 2 smaller elements (2 and 1).
To the right of 2 there is only 1 smaller element (1).
To the right of 6 there is 1 smaller element (1).
To the right of 1 there are 0 smaller elements.

Solution Ideas #

• Given an array，it is required to output, for each element in the array, the elements to the right of its position in the array that are smaller than it。
• This problem is a reduced version of Problem 327。Since it is necessary to find the elements to the right of an array position that are smaller than the current position’s element，scan from the right side of the array to the left。Construct a segment tree，where each parent node in the segment tree stores the sum of the occurrence counts of its child nodes。The given data may be very large，so discretize first when constructing the segment tree。It should also be noted that there may be duplicate elements in the array，so when constructing the segment tree, first deduplicate and sort。During the process of scanning from right to left，add the elements in the array one by one，and immediately query once after each addition。The query interval is [minNum, nums[i]-1]。If it is minNum then output 0，and also remember to insert this minimum value。The idea of this problem is generally similar to Problem 327，for details see Problem 327。
• This problem can also be solved using a Binary Indexed Tree。Compared with Problem 327, it is much simpler。The first step is still to put all used elements into the allNums array，the second step is sorting + discretization。Since the problem requires outputting the smaller elements on the right side，the third step is to insert in reverse order to construct the Binary Indexed Tree，Query queries the total number of elements in the [1,i-1] interval, which is the number of smaller elements on the right。Note that the final output is in forward order, while the computation is done in reverse order，so the answers in the final array still need to be reversed once。Similar routine problems include，Problem 327，Problem 493。

Code #

package leetcode

import (
	"sort"

	"github.com/halfrost/leetcode-go/template"
)

// Solution One Segment Tree
func countSmaller(nums []int) []int {
	if len(nums) == 0 {
		return []int{}
	}
	st, minNum, numsMap, numsArray, res := template.SegmentCountTree{}, 0, make(map[int]int, 0), []int{}, make([]int, len(nums))
	for i := 0; i < len(nums); i++ {
		numsMap[nums[i]] = nums[i]
	}
	for _, v := range numsMap {
		numsArray = append(numsArray, v)
	}
	// Sorting is to make the intervals in the segment tree satisfy left <= right，if not sorted here，many intervals in the segment tree will be invalid。
	sort.Ints(numsArray)
	minNum = numsArray[0]
	// Initialize the segment tree，the values inside the nodes are all assigned 0，that is, the count is 0
	st.Init(numsArray, func(i, j int) int {
		return 0
	})
	for i := len(nums) - 1; i >= 0; i-- {
		if nums[i] == minNum {
			res[i] = 0
			st.UpdateCount(nums[i])
			continue
		}
		st.UpdateCount(nums[i])
		res[i] = st.Query(minNum, nums[i]-1)
	}
	return res
}

// Solution Two Binary Indexed Tree
func countSmaller1(nums []int) []int {
	// copy one copy of the original array to the allNums array of all numbers
	allNums, res := make([]int, len(nums)), []int{}
	copy(allNums, nums)
	// Discretize allNums
	sort.Ints(allNums)
	k := 1
	kth := map[int]int{allNums[0]: k}
	for i := 1; i < len(allNums); i++ {
		if allNums[i] != allNums[i-1] {
			k++
			kth[allNums[i]] = k
		}
	}
	// Binary Indexed Tree Query
	bit := template.BinaryIndexedTree{}
	bit.Init(k)
	for i := len(nums) - 1; i >= 0; i-- {
		res = append(res, bit.Query(kth[nums[i]]-1))
		bit.Add(kth[nums[i]], 1)
	}
	for i := 0; i < len(res)/2; i++ {
		res[i], res[len(res)-1-i] = res[len(res)-1-i], res[i]
	}
	return res
}