diff --git a/sd.go b/sd.go
index fa46270..5f5b7aa 100644
--- a/sd.go
+++ b/sd.go
@@ -48,7 +48,7 @@ func query(v int, l int, r int, n int) (int, int) {
 		return seg[n], l
 	}
 	m := (l + r) >> 1
-	if seg[n<<1] < v {
+	if seg[n<<1] >= v {
 		return query(v, l, m, n<<1)
 	} else {
 		return query(v-seg[n<<1], m+1, r, n<<1|1)
@@ -74,6 +74,7 @@ func main() {
 		panic(err)
 	}
 	build(rows, 0, N-1, 1)
+	rows.Close()
 
 	// https://stackoverflow.com/questions/14094190/function-similar-to-getchar
 	// disable input buffering
@@ -86,16 +87,19 @@ func main() {
 	for {
 		// Choose a random card
 		x := rand.Intn(sum)
+		fmt.Println(sum)
 		fmt.Println(x)
 		w, i := query(x, 0, N-1, 1)
 		fmt.Println(w)
 		fmt.Println(i)
-
+		
+		// Get card contents from database
 		var key, val string
-		db.QueryRow("SELECT * FROM cards WHERE idx=?", i).Scan(&i, &w, &key, &val)
+		db.QueryRow("SELECT key, val FROM cards WHERE idx=?", i).Scan(&key, &val)
 		fmt.Println(key)
 		fmt.Println(val)
 
+		// Read user input
 		var b []byte = make([]byte, 1)
 		os.Stdin.Read(b)
 		if b[0] == byte('y') {
@@ -107,7 +111,12 @@ func main() {
 		} else {
 			os.Exit(0)
 		}
+
+		// Update segment tree and database
 		update(i, w, 0, N-1, 1)
-		db.Query("UPDATE cards SET weight=? WHERE idx=?", w, i)
+		_, err = db.Exec("UPDATE cards SET weight=? WHERE idx=?", w, i)
+		if err != nil {
+			panic(err)
+		}
 	}
 }