# -*- coding: iso-8859-2 -*-
#
# Trie data strucure
#
# Wojciech Muła, http://0x80.pl
# Public domain
#

# Changelog:
# 23.03.2007
#	* None become valid value (nil is used to mark empty nodes)
#	* value -> __value (priv. meth.)
# 22.03.2007
#   * fixed stupid bug in add_word
#   * added methods: __getitem__, __setitem__, get, setdefault
#   * added some tests
#  2.03.2007
#   * added 'override' switch in add_word method
# 28.02.2007
#   + __len__
#   + __iter__, iteritems, iterkeys, itervalues
#   + del_word
#   * fixed value meth.
# 27.02.2007
#   + add_word
#   + longest_prefix
#   + value
#   + __contains__

nil = object()

class Trie(object):
	def __init__(self):
		self.__next = self.a = {}
		self.__val  = nil
		self.__len  = 0


	def __contains__(self, word):
		return self.__value(word) is not nil
	

	def __len__(self):
		return self.__len


	def add_word(self, word, value, override=True):
		node = self
		for c in word:
			if c not in node.__next:
				node.__next[c] = Trie()
			node = node.__next[c]
	
		if node.__val is nil:
			self.__len += 1
			node.__val = value
		elif override: # update value for existing word
			node.__val = value


	def del_word(self, word):
		def aux(node, word, i):
			if i == len(word):
				if node.__val is not nil:
					node.__val  = nil
					self.__len -= 1
				return (not node.__next)
			else:
				try:
					c = word[i]
					if aux(node.__next[c], word, i+1):
						del node.__next[c]
						return (not node.__next and node.__val is nil)
				except KeyError:
					# word not in dict
					pass
		#aux

		aux(self, word, 0)


	def __value(self, word):
		node = self
		for c in word:
			if c in node.__next:
				node = node.__next[c]
			else:
				return nil

		return node.__val

	
	def get(self, word, default=None):
		value = self.__value(word)
		if value is not nil:
			return value
		else:
			return default


	def __setitem__(self, word, val):
		self.add_word(word, val, True)
	

	def setdefault(self, word, default):
		value = self.__value(word)
		if value is not nil:
			return value
		else:
			self.add_word(word, default)
			return default


	def __getitem__(self, word):
		value = self.__value(word)
		if value is not nil:
			return value
		else:
			raise KeyError(word)


	def longest_prefix(self, word):
		node  = self
		last  = 0
		for i, c in enumerate(word):
			if c in node.__next:
				node = node.__next[c]
				if node.__val is not nil:
					last  = i + 1
					value = node.__val
			else:
				break
		
		if last:
			return (last, value)
		else:
			return None
	

	def iteritems(self):
		stack = [(self, "")]
		while stack:
			node, s = stack.pop()
			if node.__val is not nil:
				yield s, node.__val
			for c in sorted(node.__next.keys(), reverse=True):
				stack.append( (node.__next[c], s+c) )
	
	def iterkeys(self):
		for key, val in self.iteritems():
			yield key


	def itervalues(self):
		for key, val in self.iteritems():
			yield val


	def __iter__(self):
		return self.iterkeys()

	
	def _dump(self):
		stack = [(self, "")]
		while stack:
			node, s = stack.pop()
			yield s, node.__val
			for c in sorted(node.__next.keys(), reverse=True):
				stack.append( (node.__next[c], s+c) )


if __name__ == '__main__':
	T = Trie()

	for word, value in [("trie", 123), ("triedata", 256), ("triedatastruct", 369)]:
		n = len(T)
		T.add_word(word, value)
		assert word in T
		assert T[word] == value
		assert len(T) == n+1
	
	# longest prefix, here len("triedata") = 8, T["triedata"] = 256
	assert T.longest_prefix("triedatast") == (8, 256)

	# override switch test
	new_val = 625
	T.add_word("triedata", new_val)
	assert T["triedata"] == new_val # value changed
	assert len(T) == 3
	
	T.add_word("triedata", 1024, False)
	assert T["triedata"] == new_val # value kept

	assert T.setdefault("tryit", "setdefault") == "setdefault"

	# remove test
	for word in ["triedata", "trie", "triedatastruct"]:
		n = len(T)
		T.del_word(word)
		assert word not in T
		assert len(T) == n-1
	
# vim: ts=4 sw=4 nowrap