diff options
| -rw-r--r-- | sd.go | 6 | ||||
| -rw-r--r-- | segmenttree.go | 18 |
2 files changed, 12 insertions, 12 deletions
@@ -40,7 +40,7 @@ func main() { if err != nil { panic(err) } - s.build(rows, 0, N-1, 1) + s.Build(rows, 0, N-1, 1) rows.Close() if *verbose { @@ -58,7 +58,7 @@ func main() { for { // Choose a random card x := rand.Intn(s.seg[1]) - w, i := s.query(x, 0, N-1, 1) + w, i := s.Query(x, 0, N-1, 1) if *verbose { fmt.Println(s.seg[1]) @@ -88,7 +88,7 @@ func main() { } // Update segment tree and database - s.update(i, w, 0, N-1, 1) + s.Update(i, w, 0, N-1, 1) _, err = db.Exec("UPDATE cards SET weight=? WHERE idx=?", w, i) if err != nil { panic(err) diff --git a/segmenttree.go b/segmenttree.go index e2be46a..edda5c0 100644 --- a/segmenttree.go +++ b/segmenttree.go @@ -10,42 +10,42 @@ type segmentTree struct { } // Build segment tree -func (s segmentTree) build(a *sql.Rows, l, r, n int) { +func (s segmentTree) Build(a *sql.Rows, l, r, n int) { if l == r { a.Next() a.Scan(&s.seg[n]) return } m := (l + r) >> 1 - s.build(a, l, m, n<<1) - s.build(a, m+1, r, n<<1|1) + s.Build(a, l, m, n<<1) + s.Build(a, m+1, r, n<<1|1) s.seg[n] = s.seg[n<<1] + s.seg[n<<1|1] } // Update value at index x -func (s segmentTree) update(x, v, l, r, n int) { +func (s segmentTree) Update(x, v, l, r, n int) { if l == r { s.seg[n] = v return } m := (l + r) >> 1 if x <= m { - s.update(x, v, l, m, n<<1) + s.Update(x, v, l, m, n<<1) } else { - s.update(x, v, m+1, r, n<<1|1) + s.Update(x, v, m+1, r, n<<1|1) } s.seg[n] = s.seg[n<<1] + s.seg[n<<1|1] } // Find element with prefix sum v -func (s segmentTree) query(v, l, r, n int) (int, int) { +func (s segmentTree) Query(v, l, r, n int) (int, int) { if l == r { return s.seg[n], l } m := (l + r) >> 1 if s.seg[n<<1] >= v { - return s.query(v, l, m, n<<1) + return s.Query(v, l, m, n<<1) } else { - return s.query(v-s.seg[n<<1], m+1, r, n<<1|1) + return s.Query(v-s.seg[n<<1], m+1, r, n<<1|1) } } |