Skip to content

432. All O`one Data Structure 👍

  • Time: $O(1)$
  • Space: $O(n)$
 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
struct Node {
  int count;
  unordered_set<string> keys;
};

class AllOne {
 public:
  void inc(string key) {
    if (const auto it = keyToIterator.find(key); it == keyToIterator.end())
      addNewKey(key);
    else
      incrementExistingKey(it, key);
  }

  void dec(string key) {
    const auto it = keyToIterator.find(key);
    // It is guaranteed that key exists in the data structure before the
    // decrement.
    decrementExistingKey(it, key);
  }

  string getMaxKey() {
    return nodeList.empty() ? "" : *nodeList.back().keys.begin();
  }

  string getMinKey() {
    return nodeList.empty() ? "" : *nodeList.front().keys.begin();
  }

 private:
  list<Node> nodeList;  // list of nodes sorted by count
  unordered_map<string, list<Node>::iterator> keyToIterator;

  // Adds a new node with count 1.
  void addNewKey(const string& key) {
    if (nodeList.empty() || nodeList.front().count > 1)
      nodeList.push_front({1, {key}});
    else  // nodeList.front().count == 1
      nodeList.front().keys.insert(key);
    keyToIterator[key] = nodeList.begin();
  }

  // Increments the count of the key by 1.
  void incrementExistingKey(
      unordered_map<string, list<Node>::iterator>::iterator it,
      const string& key) {
    const auto listIt = it->second;

    auto nextIt = next(listIt);
    const int newCount = listIt->count + 1;
    if (nextIt == nodeList.end() || nextIt->count > newCount)
      nextIt = nodeList.insert(nextIt, {newCount, {key}});
    else  // Node with count + 1 exists.
      nextIt->keys.insert(key);

    keyToIterator[key] = nextIt;
    remove(listIt, key);
  }

  // Decrements the count of the key by 1.
  void decrementExistingKey(
      unordered_map<string, list<Node>::iterator>::iterator it,
      const string& key) {
    const auto listIt = it->second;
    if (listIt->count == 1) {
      keyToIterator.erase(it);
      remove(listIt, key);
      return;
    }

    auto prevIt = prev(listIt);
    const int newCount = listIt->count - 1;
    if (listIt == nodeList.begin() || prevIt->count < newCount)
      prevIt = nodeList.insert(listIt, {newCount, {key}});
    else  // Node with count - 1 exists.
      prevIt->keys.insert(key);

    keyToIterator[key] = prevIt;
    remove(listIt, key);
  }

  // Removes the key from the node list.
  void remove(list<Node>::iterator it, const string& key) {
    it->keys.erase(key);
    if (it->keys.empty())
      nodeList.erase(it);
  }
};

 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class Node {
  public int count;
  public Set<String> keys = new HashSet<>();
  public Node prev = null;
  public Node next = null;
  public Node(int count) {
    this.count = count;
  }
  public Node(int count, String key) {
    this.count = count;
    keys.add(key);
  }
}

class AllOne {
  public AllOne() {
    head.next = tail;
    tail.prev = head;
  }

  public void inc(String key) {
    if (keyToNode.containsKey(key))
      incrementExistingKey(key);
    else
      addNewKey(key);
  }

  public void dec(String key) {
    // It is guaranteed that key exists in the data structure before the
    // decrement.
    decrementExistingKey(key);
  }

  public String getMaxKey() {
    return tail.prev == head ? "" : tail.prev.keys.iterator().next();
  }

  public String getMinKey() {
    return head.next == tail ? "" : head.next.keys.iterator().next();
  }

  private Map<String, Node> keyToNode = new HashMap<>();
  private Node head = new Node(0);
  private Node tail = new Node(0);

  // Adds a new node with frequency 1.
  private void addNewKey(final String key) {
    if (head.next.count == 1)
      head.next.keys.add(key);
    else
      insertAfter(head, new Node(1, key));
    keyToNode.put(key, head.next);
  }

  // Increments the frequency of the key by 1.
  private void incrementExistingKey(final String key) {
    Node node = keyToNode.get(key);
    node.keys.remove(key);

    if (node.next == tail || node.next.count > node.count + 1)
      insertAfter(node, new Node(node.count + 1));

    node.next.keys.add(key);
    keyToNode.put(key, node.next);

    if (node.keys.isEmpty())
      remove(node);
  }

  // Decrements the count of the key by 1.
  private void decrementExistingKey(final String key) {
    Node node = keyToNode.get(key);
    node.keys.remove(key);

    if (node.count > 1) {
      if (node.prev == head || node.prev.count != node.count - 1)
        insertAfter(node.prev, new Node(node.count - 1));
      node.prev.keys.add(key);
      keyToNode.put(key, node.prev);
    } else {
      keyToNode.remove(key);
    }

    if (node.keys.isEmpty())
      remove(node);
  }

  private void insertAfter(Node node, Node newNode) {
    newNode.prev = node;
    newNode.next = node.next;
    node.next.prev = newNode;
    node.next = newNode;
  }

  private void remove(Node node) {
    node.prev.next = node.next;
    node.next.prev = node.prev;
  }
}

 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from dataclasses import dataclass


@dataclass
class Node:
  def __init__(self, count: int, key: str | None = None):
    self.count = count
    self.keys: set[str] = {key} if key else set()
    self.prev: Node | None = None
    self.next: Node | None = None

  def __eq__(self, other) -> bool:
    if not isinstance(other, Node):
      return NotImplemented
    return self.count == other.count and self.keys == other.keys


class AllOne:
  def __init__(self):
    self.keyToNode: dict[str, Node] = {}
    self.head = Node(0)
    self.tail = Node(0)
    self.head.next = self.tail
    self.tail.prev = self.head

  def inc(self, key: str) -> None:
    if key in self.keyToNode:
      self._incrementExistingKey(key)
    else:
      self._addNewKey(key)

  def dec(self, key: str) -> None:
    # It is guaranteed that key exists in the data structure before the
    # decrement.
    self._decrementExistingKey(key)

  def getMaxKey(self) -> str:
    return '' if self.tail.prev == self.head \
        else next(iter(self.tail.prev.keys))

  def getMinKey(self) -> str:
    return '' if self.head.next == self.tail \
        else next(iter(self.head.next.keys))

  def _addNewKey(self, key: str) -> None:
    """Adds a new node with frequency 1."""
    if self.head.next.count == 1:
      self.head.next.keys.add(key)
    else:
      self._insertAfter(self.head, Node(1, key))
    self.keyToNode[key] = self.head.next

  def _incrementExistingKey(self, key: str) -> None:
    """Increments the frequency of the key by 1."""
    node = self.keyToNode[key]
    node.keys.remove(key)
    if node.next == self.tail or node.next.count > node.count + 1:
      self._insertAfter(node, Node(node.count + 1))
    node.next.keys.add(key)
    self.keyToNode[key] = node.next
    if not node.keys:
      self._remove(node)

  def _decrementExistingKey(self, key: str) -> None:
    """Decrements the count of the key by 1."""
    node = self.keyToNode[key]
    node.keys.remove(key)
    if node.count > 1:
      if node.prev == self.head or node.prev.count != node.count - 1:
        self._insertAfter(node.prev, Node(node.count - 1))
      node.prev.keys.add(key)
      self.keyToNode[key] = node.prev
    else:
      del self.keyToNode[key]
    if not node.keys:
      self._remove(node)

  def _insertAfter(self, node: Node, newNode: Node) -> None:
    newNode.prev = node
    newNode.next = node.next
    node.next.prev = newNode
    node.next = newNode

  def _remove(self, node: Node) -> None:
    node.prev.next = node.next
    node.next.prev = node.prev