package priorityqueue
// 参考container/heap/example_pq_test.go
// An Item is something we manage in a priority queue.
type Item struct {
value interface{} // The value of the item; arbitrary.
priority int // The priority of the item in the queue.
index int // The index of the item in the heap.
}
// A PriorityQueue implements heap.*Item and holds Items.
type PriorityQueue []*Item
// NewPriorityQueue 新建一个优先队列
func NewPriorityQueue(items ...*Item) PriorityQueue {
return items
}
func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool {
// We want Pop to give us the highest, not lowest, priority so we use greater than here.
return pq[i].priority > pq[j].priority
}
func (pq PriorityQueue) Swap(i, j int) {
pq[i], pq[j] = pq[j], pq[i]
pq[i].index = i
pq[j].index = j
}
func (pq *PriorityQueue) Push(item *Item) {
n := len(*pq)
item.index = n
*pq = append(*pq, item)
}
func (pq *PriorityQueue) Pop() *Item {
old := *pq
n := len(old)
item := old[n-1]
old[n-1] = nil // avoid memory leak
item.index = -1 // for safety
*pq = old[0 : n-1]
return item
}
// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *Item, value interface{}, priority int) {
item.value = value
item.priority = priority
Fix(pq, item.index)
}
// heap
// Init establishes the heap invariants required by the other routines in this package.
// Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// The complexity is O(n) where n = h.Len().
func Init(h *PriorityQueue) {
n := h.Len()
for i := n/2 - 1; i >= 0; i-- {
down(h, i, n)
}
}
// Push pushes the element x onto the heap.
// The complexity is O(log n) where n = h.Len().
func Push(h *PriorityQueue, x *Item) {
h.Push(x)
up(h, h.Len()-1)
}
// Pop removes and returns the minimum element (according to Less) from the heap.
// The complexity is O(log n) where n = h.Len().
// Pop is equivalent to Remove(h, 0).
func Pop(h *PriorityQueue) *Item {
n := h.Len() - 1
h.Swap(0, n)
down(h, 0, n)
return h.Pop()
}
// Remove removes and returns the element at index i from the heap.
// The complexity is O(log n) where n = h.Len().
func Remove(h *PriorityQueue, i int) *Item {
n := h.Len() - 1
if n != i {
h.Swap(i, n)
if !down(h, i, n) {
up(h, i)
}
}
return h.Pop()
}
// Fix re-establishes the heap ordering after the element at index i has changed its value.
// Changing the value of the element at index i and then calling Fix is equivalent to,
// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
// The complexity is O(log n) where n = h.Len().
func Fix(h *PriorityQueue, i int) {
if !down(h, i, h.Len()) {
up(h, i)
}
}
func up(h *PriorityQueue, j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !h.Less(j, i) {
break
}
h.Swap(i, j)
j = i
}
}
func down(h *PriorityQueue, i0, n int) bool {
i := i0
for {
j1 := 2*i + 1
if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
break
}
j := j1 // left child
if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
j = j2 // = 2*i + 2 // right child
}
if !h.Less(j, i) {
break
}
h.Swap(i, j)
i = j
}
return i > i0
}
package disjointsets
type DSU []int
// ConstructorDSU 创建容量为n的并查集
// n指的是并查集中点的个数
func ConstructorDSU(n int) DSU {
var dsu = make(DSU, n)
// 初始化时每个节点的根节点就是自身
for i := 0; i < n; i++ {
dsu[i] = i
}
return dsu
}
// Find 查找a的根节点,并进行路径压缩(注意使用路径压缩的话,独立集合的根结点不能初始化为一个相同值,比如-1)
// 如`dsu[i] = -1`,这样在合并时求解是否存在环的时候所找根节点都会为-1而造成误判
func (dsu DSU) Find(a int) int {
if dsu[a] != a {
// 某个节点的根节点不是自身时
// 路径压缩
dsu[a] = dsu.Find(dsu[a])
}
// 获取到根节点
return dsu[a]
}
// Union 合并两个集合
// False: 说明已经在一个集合中,无需合并
// True: 合并成功
func (dsu DSU) Union(a, b int) bool {
var (
aRoot = dsu.Find(a)
bRoot = dsu.Find(b)
)
// 如果根节点相同说明已经在同一个集合中了
if aRoot == bRoot {
return false
}
// 将集合a加入到集合b中
dsu[aRoot] = bRoot
return true
}
package trie
type Trie interface {
Insert(string)
Search(string) bool
HasPrefix(string) bool
}
// Type 前缀树类型
type Type int32
const (
// EmHashTrie 哈希前缀树
EmHashTrie = iota
// EmListTrie 列表前缀树
EmListTrie
// EmArrayTrie 数组前缀树
EmArrayTrie
)
// NewTrie 创建一个Trie
func NewTrie(t Type) Trie {
switch t {
case EmHashTrie:
return NewHashTrie()
case EmListTrie:
return NewListTrie()
case EmArrayTrie:
return NewArrayTrie()
default:
// HashTrie较为通用
return NewHashTrie()
}
}
package trie
// ArrayTrie 数组前缀树
type ArrayTrie struct {
isWord bool
children [26]*ArrayTrie
}
var _ Trie = (*ArrayTrie)(nil)
// NewArrayTrie new arrayTrie
func NewArrayTrie() *ArrayTrie {
return &ArrayTrie{children: [26]*ArrayTrie{}}
}
// Insert 往前缀树中添加一个元素word
func (a *ArrayTrie) Insert(word string) {
var curNode = a
for i := 0; i < len(word); i++ {
if curNode.children[word[i]-'a'] == nil {
curNode.children[word[i]-'a'] = NewArrayTrie()
}
curNode = curNode.children[word[i]-'a']
}
curNode.isWord = true
}
// Search 查找前缀树中是否元素word
func (a *ArrayTrie) Search(word string) bool {
var _, has = a.search(word)
return has
}
func (a *ArrayTrie) search(word string) (*ArrayTrie, bool) {
var curNode = a
for i := 0; i < len(word); i++ {
if curNode.children[word[i]-'a'] == nil {
return nil, false
}
curNode = curNode.children[word[i]-'a']
}
return curNode, curNode.isWord
}
// HasPrefix 查询前缀树中是否存在前缀prefix
func (a *ArrayTrie) HasPrefix(prefix string) bool {
var curNode = a
for i := 0; i < len(prefix); i++ {
if curNode.children[prefix[i]-'a'] == nil {
return false
}
curNode = curNode.children[prefix[i]-'a']
}
return true
}
package trie
// ListTrie 列表前缀树
type ListTrie struct {
isWord bool
char byte
children []*ListTrie
}
var _ Trie = (*ListTrie)(nil)
// NewListTrie new listTrie
func NewListTrie() *ListTrie {
return new(ListTrie)
}
// Insert 往前缀树中添加一个元素word
func (a *ListTrie) Insert(word string) {
var (
curNode = a
tNode *ListTrie
)
for i := 0; i < len(word); i++ {
if tNode = curNode.find(word[i]); tNode == nil {
curNode.children = append(curNode.children, &ListTrie{char: word[i]})
}
curNode = curNode.find(word[i])
}
curNode.isWord = true
}
func (a *ListTrie) find(c byte) *ListTrie {
for i := 0; i < len(a.children); i++ {
if a.children[i].char == c {
return a.children[i]
}
}
return nil
}
// Search 查找前缀树中是否元素word
func (a *ListTrie) Search(word string) bool {
var _, has = a.search(word)
return has
}
func (a *ListTrie) search(word string) (*ListTrie, bool) {
var curNode = a
for i := 0; i < len(word); i++ {
curNode = curNode.find(word[i])
if curNode == nil {
return nil, false
}
}
return curNode, curNode.isWord
}
// HasPrefix 查询前缀树中是否存在前缀prefix
func (a *ListTrie) HasPrefix(prefix string) bool {
var curNode = a
for i := 0; i < len(prefix); i++ {
curNode = curNode.find(prefix[i])
if curNode == nil {
return false
}
}
return true
}
package trie
// HashTrie Hash前缀树
type HashTrie struct {
isWord bool
children map[byte]*HashTrie
}
var _ Trie = (*HashTrie)(nil)
// NewHashTrie new hashTrie
func NewHashTrie() *HashTrie {
return &HashTrie{children: make(map[byte]*HashTrie)}
}
// Insert 往前缀树中添加一个元素word
func (h *HashTrie) Insert(word string) {
var curNode = h
for i := 0; i < len(word); i++ {
if _, has := curNode.children[word[i]]; !has {
curNode.children[word[i]] = NewHashTrie()
}
curNode = curNode.children[word[i]]
}
curNode.isWord = true
}
// Search 查找前缀树中是否元素word
func (h *HashTrie) Search(word string) bool {
var _, has = h.search(word)
return has
}
func (h *HashTrie) search(word string) (*HashTrie, bool) {
var curNode = h
for i := 0; i < len(word); i++ {
if curNode.children[word[i]] == nil {
return nil, false
}
curNode = curNode.children[word[i]]
}
return curNode, curNode.isWord
}
// HasPrefix 查询前缀树中是否存在前缀prefix
func (h *HashTrie) HasPrefix(prefix string) bool {
var curNode = h
for i := 0; i < len(prefix); i++ {
if curNode.children[prefix[i]] == nil {
return false
}
curNode = curNode.children[prefix[i]]
}
return true
}
package segmenttree
import (
"errors"
)
var (
// ErrIndexIllegal 索引非法
ErrIndexIllegal = errors.New("index is illegal")
)
package segmenttree
// Int64Merger 用户自定义区间内操作逻辑
type Int64Merger func(int64, int64) int64
type Int64Optimization func(int, int, int64) int64
// Int64SegmentTree 线段树
type Int64SegmentTree struct {
data []int64
tree []int64
lazy []int64
merge Int64Merger
optimization Int64Optimization
}
// NewInt64SegmentTree new Int64SegmentTree
func NewInt64SegmentTree(array []int64, merge Int64Merger) *Int64SegmentTree {
var seg = &Int64SegmentTree{
data: array,
tree: make([]int64, 4*len(array)),
lazy: make([]int64, 4*len(array)),
merge: merge,
}
seg.buildInt64SegmentTree(0, 0, seg.Size()-1)
return seg
}
// WithInt64Optimization 指定Int64Optimization来自定义区间[left, right]更新value的计算操作
func (s *Int64SegmentTree) WithInt64Optimization(f Int64Optimization) {
s.optimization = f
}
// leftChild 左子树index
func (s *Int64SegmentTree) leftChild(idx int) int {
return (idx << 1) + 1
}
// rightChild 右子树index
func (s *Int64SegmentTree) rightChild(idx int) int {
return (idx << 1) + 2
}
// buildInt64SegmentTree 建立线段树
func (s *Int64SegmentTree) buildInt64SegmentTree(idx int, left int, right int) {
if left == right {
s.tree[idx] = s.data[left]
return
}
var (
leftTreeIdx = s.leftChild(idx)
rightTreeIdx = s.rightChild(idx)
mid = left + (right-left)/2
)
// 创建左子树的线段树
s.buildInt64SegmentTree(leftTreeIdx, left, mid)
// 创建右子树的线段树
s.buildInt64SegmentTree(rightTreeIdx, mid+1, right)
// merge
s.tree[idx] = s.merge(s.tree[leftTreeIdx], s.tree[rightTreeIdx])
}
// Size 线段树元素数
func (s *Int64SegmentTree) Size() int {
return len(s.data)
}
// Get 索引原数组值
func (s *Int64SegmentTree) Get(idx int) (int64, error) {
if idx < 0 || idx >= s.Size() {
return 0, ErrIndexIllegal
}
return s.data[idx], nil
}
// Query 区间查询
func (s *Int64SegmentTree) Query(queryLeft int, queryRight int) (int64, error) {
if queryLeft > queryRight {
queryRight, queryLeft = queryLeft, queryRight
}
if queryLeft < 0 || queryRight >= s.Size() {
return 0, ErrIndexIllegal
}
return s.query(0, 0, s.Size()-1, queryLeft, queryRight), nil
}
func (s *Int64SegmentTree) query(idx int, left int, right int, queryLeft int, queryRight int) int64 {
if left == queryLeft && right == queryRight {
// 命中所查找区间
return s.tree[idx]
}
var (
mid = left + (right-left)/2
// 计算左右子节点下标
leftIdx = s.leftChild(idx)
rightIdx = s.rightChild(idx)
)
if queryLeft >= mid+1 {
// 所查询区间在右半部分
return s.query(rightIdx, mid+1, right, queryLeft, queryRight)
}
if queryRight <= mid {
// 所查询区间在左半部分
return s.query(leftIdx, left, mid, queryLeft, queryRight)
}
// 所查询区间分布在左右区间
return s.merge(
// 查询左部分
s.query(leftIdx, left, mid, queryLeft, mid),
// 查询右部分
s.query(rightIdx, mid+1, right, mid+1, queryRight),
)
}
// QueryLazy 懒惰查询
func (s *Int64SegmentTree) QueryLazy(queryLeft int, queryRight int) (int64, error) {
if queryLeft > queryRight {
queryRight, queryLeft = queryLeft, queryRight
}
if queryLeft < 0 || queryRight >= s.Size() {
return 0, ErrIndexIllegal
}
return s.queryLazy(0, 0, s.Size()-1, queryLeft, queryRight), nil
}
func (s *Int64SegmentTree) queryLazy(idx int, left int, right int, queryLeft int, queryRight int) int64 {
// 处理懒惰更新
s.pushDown(idx, left, right)
var (
mid = left + (right-left)>>1
leftIdx = s.leftChild(idx)
rightIdx = s.rightChild(idx)
)
if left > right || left > queryRight || right < queryLeft {
return 0
}
if queryLeft <= left && right <= queryRight {
// 在所查找区间范围内
return s.tree[idx]
}
if queryLeft >= mid+1 {
// 所查询区间在右半部分
return s.queryLazy(rightIdx, mid+1, right, queryLeft, queryRight)
}
if queryRight <= mid {
// 所查询区间在左半部分
return s.queryLazy(leftIdx, left, mid, queryLeft, queryRight)
}
// 所查询区间分布在左右区间
return s.merge(
// 查询左部分
s.queryLazy(leftIdx, left, mid, queryLeft, mid),
// 查询右部分
s.queryLazy(rightIdx, mid+1, right, mid+1, queryRight),
)
}
// Set 更新元素值
func (s *Int64SegmentTree) Set(index int, e int64) error {
if index < 0 || index >= s.Size() {
return ErrIndexIllegal
}
// 更新数组元素值
s.data[index] = e
// 更新tree元素值
s.set(0, 0, s.Size()-1, index, e)
return nil
}
func (s *Int64SegmentTree) set(idx int, left int, right int, index int, e int64) {
if left == right {
// 命中节点,更新元素值
s.tree[idx] = e
return
}
var (
leftIdx = s.leftChild(idx)
rightIdx = s.rightChild(idx)
mid = left + (right-left)/2
)
if index <= mid {
// idx在左边
s.set(leftIdx, left, mid, index, e)
} else {
// idx在右边
s.set(rightIdx, mid+1, right, index, e)
}
// merge
s.tree[idx] = s.merge(s.tree[leftIdx], s.tree[rightIdx])
}
// AddValueLazy
// 给[addLeft....addRight]位置的值都加上value
// 注意这里的更新值是在原来值的基础上增加或者减少,而不是把这个区间内的值都赋值为 x,区间更新和单点更新不同
// 这里的区间更新关注的是变化,单点更新关注的是定值
// 当然区间更新也可以都更新成定值,如果只区间更新成定值,那么 lazy 更新策略需要变化,merge 策略也需要变化,这里暂不详细讨论
func (s *Int64SegmentTree) AddValueLazy(addLeft int, addRight int, value int64) error {
if addLeft > addRight {
addRight, addLeft = addLeft, addRight
}
if addLeft < 0 || addRight >= s.Size() {
return ErrIndexIllegal
}
s.addValueLazy(0, 0, s.Size()-1, addLeft, addRight, value)
return nil
}
func (s *Int64SegmentTree) addValueLazy(idx int, left int, right int, addLeft int, addRight int, value int64) {
// 处理懒惰更新
s.pushDown(idx, left, right)
var (
mid = left + (right-left)>>1
leftIdx = s.leftChild(idx)
rightIdx = s.rightChild(idx)
)
if left > right || left > addRight || right < addLeft {
return
}
if addLeft <= left && right <= addRight {
// 正好在一个区间内区间
if s.optimization != nil {
s.tree[idx] = s.merge(s.tree[idx], s.optimization(left, right, value))
} else {
for i := 0; i < right-left+1; i++ {
s.tree[idx] = s.merge(s.tree[idx], value)
}
}
if left != right {
s.lazy[leftIdx] = s.merge(s.lazy[leftIdx], value)
s.lazy[rightIdx] = s.merge(s.lazy[rightIdx], value)
}
return
}
// 需要分别更新左右区间
s.addValueLazy(leftIdx, left, mid, addLeft, addRight, value)
s.addValueLazy(rightIdx, mid+1, right, addLeft, addRight, value)
s.tree[idx] = s.merge(s.tree[leftIdx], s.tree[rightIdx])
}
func (s *Int64SegmentTree) pushDown(idx int, left int, right int) {
var (
leftIdx = s.leftChild(idx)
rightIdx = s.rightChild(idx)
)
// 处理懒惰更新
if s.lazy[idx] != 0 {
if s.optimization != nil {
s.tree[idx] = s.merge(s.tree[idx], s.optimization(left, right, s.lazy[idx]))
} else {
// 懒惰更新根节点 如果是区间和则等同于:s.tree[idx] += (right-left+1)*s.lazy[idx]
for i := 0; i < right-left+1; i++ {
s.tree[idx] = s.merge(s.tree[idx], s.lazy[idx])
}
}
if left != right {
// 懒惰更新子节点 如果是区间求和则等同于:s.lazy[leftIdx] += s.lazy[idx]
s.lazy[leftIdx] = s.merge(s.lazy[leftIdx], s.lazy[idx])
s.lazy[rightIdx] = s.merge(s.lazy[rightIdx], s.lazy[idx])
}
// 消除懒惰更新标志
s.lazy[idx] = 0
}
}