From fd72748b00e221ab4523ee104af045622d954f07 Mon Sep 17 00:00:00 2001
From: pymia <mia5419@gmail.com>
Date: Wed, 13 Sep 2017 20:34:29 +0800
Subject: [PATCH] add discard and remove function

---
 pygtrie.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/pygtrie.py b/pygtrie.py
index 0c41e05..6258ee5 100644
--- a/pygtrie.py
+++ b/pygtrie.py
@@ -1364,13 +1364,35 @@ def add(self, key):
             self._trie[key:] = True
 
     def discard(self, key):
-        raise NotImplementedError(
-            'Removing keys from PrefixSet is not implemented.')
+        try:
+            node, _ = self._get_node(key)
+        except:
+            pass
+
+        if node and node.children:
+            for child in node.children:
+                self.discard(child)
+        else:
+            node.children = {}
+            node.value = _SENTINEL
+        
+
 
+    # Raises KeyError if elem is not contained in the set.
     def remove(self, key):
-        raise NotImplementedError(
-            'Removing keys from PrefixSet is not implemented.')
+        try:
+            node, _ = self._get_node(key)
+        except KeyError:
+            return 0
+
+        if node and node.children:
+            for child in node.children:
+                self.discard(child)
+        else:
+            node.children = {}
+            node.value = _SENTINEL
 
     def pop(self):
+        
         raise NotImplementedError(
             'Removing keys from PrefixSet is not implemented.')