qs-partition.py
2.59 KB
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
47
48
49
50
51
52
53
54
55
56
57
from codanim.data import Array, Struct, Value, Box
from codanim.flow import FUNC, BLOCK, ENV, WS, DECL, WHILE, EXPR, DO, STMT, RETURN, IF
a = Array([3, 13, 5, 0, 7, 11, 4, 9, 14, 2, 10, 1, 12, 8, 6],
index=["a", "b"])
t = Struct({"len": Value(15), "val": a.ptr()})
partition = FUNC(BLOCK(ENV("t", t),
WS(" "),
DECL("pivot", EXPR("t.val[(t.len-1)//2]",
src="t.val[(t.len-1)/2]"),
animate=True,
src="int {name} = {init};"),
WS("\n "),
DECL("a", EXPR("-1"),
animate=True,
src="uint {name} = {init};"),
WS("\n "),
DECL("b", EXPR("t.len"),
animate=True,
src="uint {name} = {init};"),
WS("\n "),
WHILE(EXPR("1", src="1"),
BLOCK(WS(" "),
DO(STMT("a+=1", src="a++;"),
EXPR("t.val[a] < pivot"),
src="do {{ {body} }} while ({cond});"),
WS("\n "),
DO(STMT("b-=1", src="b--;"),
EXPR("t.val[b] > pivot"),
src="do {{ {body} }} while ({cond});"),
WS("\n "),
IF(EXPR("a >= b"),
RETURN(EXPR("b"),
src="return {value};"),
src="if ({cond}) {{ {then} }}"),
WS("\n "),
STMT("_old = t.val[a], t.val[b]",
"t.val[b], t.val[a] = _old",
src="swap(t, a, b);"),
WS("\n")),
src="while ({cond}) {{\n{body} }}"),
WS("\n")),
src="uint partition (Tab t) {{\n{body}}}\n")
partition.IP += 1 # do not highlight first step
partition() # simulate execution
# save code animation
with open("out.code", "w") as out :
out.write(partition.tex())
# save data animation
b = Box(a, t, grow="left")
with open("out.tikz", "w") as out :
out.write(b.tex(tikzpicture={"scale": .6},
tail=r"\node at (0,-2) {};"))