def build_trie(n, genes, health):
trie = Trie()
for i in range(n):
gene_trie = trie.add(genes[i])
add_health_to_gene_trie(gene_trie, health[i], i)
return trie
def test_add_health_to_gene_trie_when_add_to_beginning():
trie = Trie('a')
trie.value = [DataItem(2, 8)]
index = 1
health = 4
add_health_to_gene_trie(trie, health, index)
result = trie.value
assert len(result) == 2
assert result[0].index == 1
assert result[0].sum == 4
assert result[1].index == 2
assert result[1].sum == 12
def test_when_inserting_all_nodes_a_node_should_have_len_2():
trie = Trie()
trie.add(test_str_1)
trie.add(test_str_2)
trie.add(test_str_3)
trie.add(test_str_4)
trie.add(test_str_5)
trie.add(test_str_6)
a_node = trie.find_node('a')
assert len(a_node.children) == 2
def test_should_return_same_inserted_values():
trie = Trie()
try_test_value(trie, test_str_1)
try_test_value(trie, test_str_2)
try_test_value(trie, test_str_3)
try_test_value(trie, test_str_4)
try_test_value(trie, test_str_5)
try_test_value(trie, test_str_6)
def test_add_health_to_gene_trie_when_add_to_end():
trie = Trie('a')
trie.value = [DataItem(2, 8), DataItem(5, 14)]
index = 7
health = 4
add_health_to_gene_trie(trie, health, index)
result = trie.value
assert len(result) == 3
assert result[0].index == 2
assert result[0].sum == 8
assert result[1].index == 5
assert result[1].sum == 14
assert result[2].index == 7
assert result[2].sum == 18
def count_health(genes_data, firstIndex, lastIndex, stream):
total_health = 0
genes_count_cache = Trie()
for i in range(firstIndex, lastIndex + 1):
gene = genes_data.genes[i]
health = genes_data.health[i]
gene_count_cache_node = genes_count_cache.find_node(gene)
if gene_count_cache_node == None:
prefix_array = genes_data.prefix_arrays[i]
appear_positions = kmp(gene, stream, prefix_array)
gene_count_cache_node = genes_count_cache.add(gene)
gene_count_cache_node.value = len(appear_positions)
total_health += gene_count_cache_node.value * health
return total_health
def test_when_not_found_should_return_none():
trie = Trie()
trie.add(test_str_1)
trie.add(test_str_2)
trie.add(test_str_3)
testNode = trie.find_node(test_str_not_found)
assert testNode == None
def test_add_health_to_gene_trie_when_empty_should_add_new_data_item():
trie = Trie('a')
index = 2
health = 13
add_health_to_gene_trie(trie, health, index)
result = trie.value
assert len(result) == 1
assert result[0].index == index
assert result[0].sum == health
def test_add_health_to_gene_trie_when_inserting_in_middle_should_insert_in_correct_position_and_update_forward_sum():
trie = Trie('a')
trie.value = [DataItem(2, 8), DataItem(5, 14), DataItem(6, 19)]
index = 3
health = 4
add_health_to_gene_trie(trie, health, index)
result = trie.value
assert len(result) == 4
assert result[0].index == 2
assert result[0].sum == 8
assert result[1].index == 3
assert result[1].sum == 12
assert result[2].index == 5
assert result[2].sum == 18
assert result[3].index == 6
assert result[3].sum == 23
def test_when_empty_string_should_return_root():
trie = Trie()
testNode = trie.find_node('')
assert not testNode == None
def test_find_prefix(self):
"""Test for find_prefid
Args:
self: TestTrie
Returns:
None
Raises:
None
"""
# Given
trie_obj = Trie()
# When
trie_obj.add("hello")
trie_obj.add("hell")
trie_obj.add("he")
trie_obj.add("she")
# Then
self.assertEqual(trie_obj.find_prefix("he"), (True, 3))
self.assertEqual(trie_obj.find_prefix("hell"), (True, 2))
self.assertEqual(trie_obj.find_prefix("hello"), (True, 1))
self.assertEqual(trie_obj.find_prefix("h"), (True, 3))
self.assertEqual(trie_obj.find_prefix("me"), (False, 0))
self.genes = genes
self.health = health
self.prefix_arrays = prefix_arrays
#!/bin/python
if __name__ == '__main__':
n = int(raw_input())
genes = raw_input().rstrip().split()
health = map(int, raw_input().rstrip().split())
s = int(raw_input())
genes_prefix_cache_trie = Trie()
prefix_arrays = [None for g in genes]
# calculate the prefix arrays by using a trie as cache
for i in range(len(genes)):
gene = genes[i]
cache_trie_node = genes_prefix_cache_trie.find_node(gene)
if cache_trie_node != None:
prefix_arrays[i] = cache_trie_node.value
else:
prefix_arrays[i] = kmp_initialize_suffix_prefix_array(gene)
cache_trie_node = genes_prefix_cache_trie.add(gene)
cache_trie_node.value = prefix_arrays[i]
genes_data = GenesData(genes, health, prefix_arrays)