golang刷leetcode技巧之如何实现数字流的秩

如何实现数字流的秩

package main

import (
	"fmt"
)

type TreeNode struct {
	Val         int
	Left, Right *TreeNode
	Rank        int // 存储节点在树中的秩
}

type StreamRank struct {
	Root *TreeNode
}

// 插入节点
func (sr *StreamRank) Insert(val int) {
	if sr.Root == nil {
		sr.Root = &TreeNode{Val: val, Rank: 1}
	} else {
		sr.insertNode(sr.Root, val)
	}
}

// 递归插入节点
func (sr *StreamRank) insertNode(node *TreeNode, val int) {
	if val <= node.Val {
		node.Rank++
		if node.Left == nil {
			node.Left = &TreeNode{Val: val, Rank: 1}
		} else {
			sr.insertNode(node.Left, val)
		}
	} else {
		if node.Right == nil {
			node.Right = &TreeNode{Val: val, Rank: 1}
		} else {
			sr.insertNode(node.Right, val)
		}
	}
}

// 获取数字流中小于或等于给定值的节点数量
func (sr *StreamRank) GetRankOfNumber(val int) int {
	return sr.getRank(sr.Root, val)
}

// 递归获取节点数量
func (sr *StreamRank) getRank(node *TreeNode, val int) int {
	if node == nil {
		return 0
	}
	if val == node.Val {
		return node.Rank
	}
	if val < node.Val {
		return sr.getRank(node.Left, val)
	}
	return node.Rank + sr.getRank(node.Right, val)
}

func main() {
	sr := &StreamRank{}
	sr.Insert(5)
	sr.Insert(1)
	sr.Insert(4)
	sr.Insert(4)
	sr.Insert(5)
	sr.Insert(9)
	sr.Insert(7)
	sr.Insert(13)
	sr.Insert(3)

	fmt.Println(sr.GetRankOfNumber(1))   // 0
	fmt.Println(sr.GetRankOfNumber(3))   // 1
	fmt.Println(sr.GetRankOfNumber(4))   // 3
	fmt.Println(sr.GetRankOfNumber(9))   // 6
	fmt.Println(sr.GetRankOfNumber(13))  // 9
	fmt.Println(sr.GetRankOfNumber(100)) // 9
}

1. 思路：

要实现数字流的秩，可以使用二叉搜索树的数据结构：每个节点存储一个值和该值在树中的秩（节点左子树节点个数+1），左子树存储小于或等于该节点值的数，右子树存储大于该节点值的数。插入节点时，根据节点大小决定插入位置，并更新每个节点的秩。获取秩时，通过递归遍历树，根据节点值和目标值的大小关系，递归地定位到目标值所在的节点，并返回节点的秩。

2. 实现：

首先定义一个二叉树节点结构TreeNode，包含值Val、左子树Left、右子树Right和秩Rank。再定义一个数字流秩结构StreamRank，包含根节点Root。插入节点时，如果根节点为空，直接插入；否则调用私有方法insertNode进行递归插入节点。递归插入节点时，根据节点值和目标值的大小关系，更新节点的秩，并确定插入位置。获取秩时，调用私有方法getRank进行递归遍历树，直到找到目标值所在的节点或到达叶子节点为止，并返回节点的秩。

3. 示例：

首先创建一个StreamRank对象sr。然后按顺序插入数字5、1、4、4、5、9、7、13、3。最后分别获取数字流中小于或等于1、3、4、9、13、100的节点数量，输出结果为0、1、3、6、9、9。