update

hard/432.all-oone-data-structure.go

@@ -1,6 +1,6 @@
 type none struct{}
 
-type pair {
+type pair struct {
 	key int
 	set map[string]none
 }
@@ -22,16 +22,21 @@ func Constructor() AllOne {
 func (this *AllOne) Inc(key string) {
 	if val, ok := this.m[key]; ok { // Increments
 		next := val.Next()
-		if next == nil || val.(pair).key+1 != next.(pair).key { // Insert after
-			next = 
+		if next == nil || val.Value.(pair).key+1 < next.Value.(pair).key { // Insert after
+			next = this.li.InsertAfter(pair{val.Value.(pair).key + 1, map[string]none{key: none{}}}, val)
 		} else { // Increment after
-			
+			next.Value.(pair).set[key] = none{}
 		}
 		this.m[key] = next
+		if len(val.Value.(pair).set) == 1 { // Remove current element
+			this.li.Remove(val)
+		} else { // Remove from set
+			delete(val.Value.(pair).set, key)
+		}
 	} else { // Inserts
-		if this.li.Len() == 0 || this.li.Front().Value.(pair).key != 1 {
-			this.li.PushFront(pair{1, map[string]none{key: none{}})
-		} else {
+		if this.li.Len() == 0 || 1 < this.li.Front().Value.(pair).key { // Push front
+			this.li.PushFront(pair{1, map[string]none{key: none{}}})
+		} else { // Move to ones set
 			this.li.Front().Value.(pair).set[key] = none{}
 		}
 		this.m[key] = this.li.Front()
@@ -41,30 +46,50 @@ func (this *AllOne) Inc(key string) {
 /** Decrements an existing key by 1. If Key's value is 1, remove it from the data structure. */
 func (this *AllOne) Dec(key string) {
 	if val, ok := this.m[key]; ok {
-		
+		if val.Value.(pair).key == 1 { // Removes
+			if len(val.Value.(pair).set) == 1 { // Remove whole element
+				this.li.Remove(val)
+			} else { // Remove from set
+				delete(val.Value.(pair).set, key)
+			}
+			delete(this.m, key) // Remove from map
+		} else { // Decrements
+			prev := val.Prev()
+			if prev == nil || prev.Value.(pair).key < val.Value.(pair).key-1 { // Insert before
+				prev = this.li.InsertBefore(pair{val.Value.(pair).key - 1, map[string]none{key: none{}}}, val)
+			} else { // Move to prev set
+				prev.Value.(pair).set[key] = none{}
+			}
+			this.m[key] = prev
+			if len(val.Value.(pair).set) == 1 { // Remove current element
+				this.li.Remove(val)
+			} else {
+				delete(val.Value.(pair).set, key) // Remove from set
+			}
+		}
 	}
 }
 
 /** Returns one of the keys with maximal value. */
 func (this *AllOne) GetMaxKey() string {
-	if this.li.Len() == 0 {
-		return ""
-	}
-	set := this.li.Back().Value.(pair).set
-	for k := range set {
-		return k
+	if this.li.Len() != 0 {
+		set := this.li.Back().Value.(pair).set
+		for k := range set {
+			return k
+		}
 	}
+	return ""
 }
 
 /** Returns one of the keys with Minimal value. */
 func (this *AllOne) GetMinKey() string {
-	if this.li.Len() == 0 {
-		return ""
-	}
-	set := this.li.Front().Value.(pair).set
-	for k := range set {
-		return k
+	if this.li.Len() != 0 {
+		set := this.li.Front().Value.(pair).set
+		for k := range set {
+			return k
+		}
 	}
+	return ""
 }
 
 /**