数据结构

双端队列

可从队列2端进行添加和移除数据，效率比list高。

1
2
3
4
5
6
7
8
9
from collections import deque

d=deque([1,2,3])  # 使用list初始化
d.append(1)       # 队尾添加元素
d.appendleft(1)   # 队头添加元素
d.pop()           # 队尾移除元素
d.popleft()       # 队头移除元素
d.extend([4,5,6])        # 队尾添加多个元素
d.extendleft([4,5,6])    # 队头添加多个元素

优先队列

基于堆实现的优先队列，默认小顶堆，最小的元素在堆顶。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from queue import PriorityQueue

q = PriorityQueue()

q.put((2, 'code'))
q.put((1, 'eat'))
q.put((3, 'sleep'))

while not q.empty():
    next_item = q.get()
    print(next_item)

# 结果：
#   (1, 'eat')
#   (2, 'code')
#   (3, 'sleep')

PriorityQueue内部使用heapq函数实现，将基于函数的接口封装为了基于类的接口，同时PriorityQueue是同步的，提供了锁语义来支持多个并发的生产者和消费者。

函数式heapq

默认小顶堆

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import heapq

q = []

heapq.heappush(q, (2, 'code'))
heapq.heappush(q, (1, 'eat'))
heapq.heappush(q, (3, 'sleep'))

while q:
    next_item = heapq.heappop(q)
    print(next_item)

# 结果：
#   (1, 'eat')
#   (2, 'code')
#   (3, 'sleep')

# 将一个数组初始化为堆结构
heapq.heapify([3,2,1])

默认初始化字典

如果获取的键不存在，则采用指定的构造函数为这个键设置一个默认值。

1
2
3
4
5
6
7
8
9
10
11
from collections import defaultdict

s = 'mississippi'
d = defaultdict(int)
for k in s:
    d[k] += 1
print(d)   # defaultdict(<class 'int'>, {'m': 1, 'i': 4, 's': 4, 'p': 2})

d = defaultdict(list)
d['a'].append(1)
print(d)   # defaultdict(<class 'list'>, {'a': [1]})

顺序存储字典

按顺序存储字典键值，内部由普通字典和双向链表实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from collections import OrderedDict

user_dict = OrderedDict()
user_dict['b'] = '0xbug'
user_dict['a'] = '0bug'
user_dict['c'] = '1bug'

# popitem 弹出最后一个组
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug'), ('c', '1bug')])
print(user_dict.popitem())  # ('c', '1bug')
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug')])

# popitem 弹出第一个组
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug'), ('c', '1bug')])
print(user_dict.popitem(last=False))  # ('b', '0xbug')
print(user_dict)  # OrderedDict([('a', '0bug'), ('c', '1bug')])

# pop弹出指定元素，必须写参数
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug'), ('c', '1bug')])
print(user_dict.pop('a'))  # 0bug
print(user_dict)           # OrderedDict([('b', '0xbug'), ('c', '1bug')])

# move_to_end将指定元素移至最后
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug'), ('c', '1bug')])
print(user_dict.move_to_end('a'))  # 0bug
print(user_dict)  # OrderedDict([('b', '0xbug'), ('c', '1bug'), ('a', '0bug')])

# 获取第一个值
print(user_dict)  # OrderedDict([('b', '0xbug'), ('a', '0bug'), ('c', '1bug')])
print(next(iter(user_dict.values())))  # 0xbug

有序集合

集合元素是有序的，添加元素后集合仍然保持有序。

1
2
3
4
5
6
7
from sortedcontainers import SortedList

sl = SortedList()
sl.add(num)
sl.pop(0) # remove index 0
sl.pop()  # equals s1.pop(-1)
s1.remove(num)

键有序字典

字典的键是有序的，可按顺序遍历键。

1
2
3
4
from sortedcontainers import SortedDict

sd = SortedDict({'c': 3, 'a': 1, 'b': 2})
print(sd)    #SortedDict({'a': 1, 'b': 2, 'c': 3})

sortedcontainers非python默认库，需进行安装pip install sortedcontainers。

计数器

一种特殊的默认初始化字典，值是int类型，表示键的数量。

1
2
3
4
5
from collections import Counter

obj = Counter('aabbccc')
obj['d'] += 1
print(obj)  # Counter({'c': 3, 'a': 2, 'b': 2, 'd': 1})

自定义数据结构

并查集

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class UnionFind:
    def __init__(self,n):
        self.un=list(range(n))
        
    def find(self,p):
        if self.un[p]!=p:
            self.un[p]=self.find(self.un[p])
        return self.un[p]
    
    def merge(self,p,q):
        pr=self.find(p)
        qr=self.find(q)
        if pr!=qr:
            self.un[pr]=qr

    def connect(self,p,q):
        return self.find(p)==self.find(q)

单词树

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Trie:
    def __init__(self):
        self.children={}
        self.cnt=0
    
    def add(self,s):
        curr=self
        for c in s:
            if not c in curr.children:
                curr.children[c]=Trie()
            curr=curr.children[c]
        curr.cnt+=1

    def get(self,s):
        curr=self
        for c in s:
            if not c in curr.children:
                return 0
            curr=curr.children[c]
        return curr.cnt

树状数组

详细介绍参见：数据结构之树状数组

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class BinTree:
    def __init__(self,nums):
        self.nums=nums
        self.tree=[0]*(len(nums)+1)
        for i,num in enumerate(nums):
            self.add(i + 1,num)

    def add(self,i,v):
        if i >= len(self.tree):
            return
        self.tree[i]+=v
        self.add(i + self.lowbit(i),v)

    def ask(self,i):
        if i == 0:
            return 0
        return self.tree[i] + self.ask(i - self.lowbit(i))

    def lowbit(self,num):
        return num & (-num)
    
    def update(self, index: int, val: int) -> None:
        diff=val-self.nums[index]
        self.nums[index]=val        # 更新原数组
        self.add(index + 1,diff)    # 更新树状数组

    def sumRange(self, left: int, right: int) -> int:
        return self.ask(right+1)-self.ask(left)

线段树

详细介绍参见：数据结构之线段树，以下为标准版实现代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class Node:
    def __init__(self, l, r,v=0):
        self.left = None
        self.right = None
        self.l = l
        self.r = r
        self.m = (l + r) >> 1
        self.v = v


class SegmentTree:
    def __init__(self,nums):
        self.nums = nums
        self.root = self.__build(0,len(nums) - 1)

    def __build(self,l,r):
        if l == r:            
            return Node(l,r,self.nums[l])
        node = Node(l,r)
        node.left = self.__build(node.l,node.m)
        node.right = self.__build(node.m + 1,node.r)
        node.v = node.left.v + node.right.v
        return node

    def __modify(self, l, r, v, node):
        if l > node.r or r < node.l:
            return
        if l <= node.l and node.r <= r:
            node.v = v
            return
        self.__modify(l, r, v, node.left)
        self.__modify(l, r, v, node.right)
        node.v = node.left.v + node.right.v

    def __query(self, l, r,node):
        if l > node.r or r < node.l:
            return 0
        if l <= node.l and node.r <= r:
            return node.v
        return self.__query(l, r, node.left) + self.__query(l, r, node.right)
                
    def modify(self,index,val):
        self.__modify(index,index,val,self.root)

    def query(self,left,right):
        return self.__query(left,right,self.root)

二分模块

通过二分算法快速在有序集合中找到插入点。

1
2
3
4
5
6
7
8
9
10
11
12
13
import bisect

# 插入操作
bisect.insort_left([1,2],1)  # 插入在左侧，无返回值
bisect.insort_right([1,2],1) # 插入在右侧，无返回值

# 查找操作
bisect.bisect_left([1,2],1)  # 不插入值，返回左侧的索引，输出：0
bisect.bisect_right([1,2],1) # 不插入值，返回右侧的索引，输出：1

# 下面两行就是insort_right和bisect_right的别名
bisect.insort([1,2],1) # 插入在右侧
bisect.bisect([1,2],1) # 不插入值，返回右侧的索引

随机数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import random

# 浮点数[0,1)
random.random()

# 浮点数[2,6]
random.uniform(2, 6)

# 整数[6,8]
random.randint(6,8)

# 随机返回range中的一个元素
random.randrange(0, 31, 2)   # 输出[0,30]内的偶数

# 随机返回集合中的一个元素
random.choice([1,2,3])

# 随机打乱集合，无返回
arr=[1,2,3,4,5]
random.shuffle(arr)
print(arr)

# 从列表中随机返回3个元素，作为新列表返回
random.sample([1,2,3,4,5], 3)

累积运算

1
2
3
4
5
6
7
8
9
10
11
12
import itertools
import operator
 
if __name__ == '__main__':
    data = [1, 2, 3, 4, 5]
    # 计算前缀和
    print(list(itertools.accumulate(data)))     # [1,3,6,10,15]
    # 计算到当前位置累积相乘得结果
    data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
    print(list(itertools.accumulate(data, operator.mul, initial=2)))  # [2,6,24,144,288,288,2592,0,0,0,0]
    # 计算到当前位置的最大值并且输出
    print(list(itertools.accumulate(data, max)))  # [3,4,6,6,6,9,9,9,9,9]

排列组合

笛卡尔积

有放回抽样排列

1
2
3
4
5
6
import itertools
for i in itertools.product('ABCD', repeat = 2):
    print(i)
# 输出
# ('A', 'A') ('A', 'B') ('A', 'C') ('A', 'D') ('B', 'A') ('B', 'B') ('B', 'C') ('B', 'D') 
# ('C', 'A') ('C', 'B') ('C', 'C') ('C', 'D') ('D', 'A') ('D', 'B') ('D', 'C') ('D', 'D')

数量：$ n^r=4^2=16 $

排列

不放回抽样排列

1
2
3
4
5
6
import itertools
for i in itertools.permutations('ABCD', 2):
    print(i)
# 输出
# ('A', 'B') ('A', 'C') ('A', 'D') ('B', 'A') ('B', 'C') ('B', 'D') ('C', 'A') ('C', 'B') 
# ('C', 'D') ('D', 'A') ('D', 'B') ('D', 'C')

数量：$ A_n^r=A_4^2=\frac{4!}{(4-2)!}=12 $

组合，没有重复

不放回抽样组合

1
2
3
4
5
import itertools
for i in itertools.combinations('ABCD', 2):
    print(i)
# 输出
# ('A', 'B') ('A', 'C') ('A', 'D') ('B', 'C') ('B', 'D') ('C', 'D')

数量： $ C_n^r=C_4^2=\frac{4!}{(4-2)!*2!}=6 $

组合，有重复

有放回抽样组合

1
2
3
4
5
6
import itertools
for i in itertools.combinations_with_replacement('ABCD', 2):
    print(i)
# 输出
# ('A', 'A') ('A', 'B') ('A', 'C') ('A', 'D') ('B', 'B') ('B', 'C') ('B', 'D') ('C', 'C') 
# ('C', 'D') ('D', 'D')

相当于多增加(r-1)个元素参与选择，这(r-1)个元素是替换符，可以替换成组合中已存在的元素，形成重复。

数量： $ C_{n+r-1}^r=C_5^2=\frac{5!}{(5-2)!*2!}=10 $

相比itertools.combinations('ABCD', 2)，多出的4个组合是('A', 'X') ('B', 'X') ('C', 'X') ('D', 'X')，对应的则是('A', 'A') ('B', 'B') ('C', 'C') ('D', 'D')。