Secret-Group-Code/SubSegTree.hs at master · divanik/Secret-Group-Code · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
module SubSegTree where

data Node  = Empty | Leaf {size :: Integer} | Mid {free :: Integer, size :: Integer, lef :: Node, rig :: Node} deriving Show

make :: Integer -> Node
make s  	= Leaf s

size0 :: Node -> Integer
size0 Empty = 1
size0 all = size all

free0 :: Node -> Integer
free0 Empty = 0
free0 (Leaf sz) = sz
free0 all = free all


push :: Integer -> Node -> (Maybe Integer, Node)
push k Empty = (Nothing, Empty)

push k all@(Leaf sz) =
	if ((k < 0) || (k >= sz)) then
		(Nothing, all)
	else
		if (sz == 1) then
			(Just 0, Empty)
		else
			let a = div sz 2 in
			let b = div (sz + 1) 2 in
			if (k < a) then
				let (Just num, tr) = (push k (make a)) in
					(Just num, Mid (sz - 1) sz tr (make b))
			else
				let (Just num, tr) = (push (k - a) (make b)) in
					(Just (num + a), Mid (sz - 1) sz (make a) tr)

push k all@(Mid fr sz lf rg) =
	if ((k < 0) || (k >= fr)) then
		(Nothing, all)
	else
		let lfr = free0 lf in
		let lsz = size0 lf in
		if (k < lfr) then
			let (Just num, tr) = (push k lf) in
			(Just num, Mid (fr - 1) sz tr rg)
		else
			let (Just num, tr) = (push (k - lfr) rg) in
			(Just (num + lsz), Mid (fr - 1) sz lf tr)


push2 :: Integer -> Node -> (Maybe Integer, Node)
push2 k Empty = push k Empty

push2 k all@(Leaf sz) = push k all

push2 k all@(Mid fr sz lf rg) =
	if ((k < 0) || (k >= sz)) then
		(Nothing, all)
	else
		let lfr = free0 lf in
		let lsz = size0 lf in
		if (k < lsz) then
			let (lget, tr) = (push2 k lf) in
			case lget of
			Nothing  -> (Nothing, all)
			Just num -> (Just num, Mid (fr - 1) sz tr rg)
		else
			let (rget, tr) = (push2 (k - lsz) rg) in
			case rget of
			Nothing  -> (Nothing, all)
			Just num -> (Just (num + lsz), Mid (fr - 1) sz lf tr)