677.键值映射

题目描述

键值映射

思路：扫描、前缀树

扫描

将键值对进行存储，搜索给定的前缀prefix时，依次搜索所有键值，如果键值包含给定的前缀，则将其val进行相加，返回所有符合要求的val的和。

Python实现

# 暴力扫描
class MapSum:

    def __init__(self):
        self.map = {}


    def insert(self, key: str, val: int) -> None:
        self.map[key] = val


    def sum(self, prefix: str) -> int:
        res = 0
        for key, val in self.map.items():
            if key.startswith(prefix):
                res += val
        return res

Java实现

class MapSum {
    Map<String, Integer> map;

    public MapSum() {
        map = new HashMap<>();
    }
    
    public void insert(String key, int val) {
        map.put(key,val);
    }
    
    public int sum(String prefix) {
        int res = 0;
        for (String s: map.keySet()) {
            if (s.startsWith(prefix)) {
                res += map.get(s);
            }
        }
        return res;
    }
}

前缀树

由于要处理前缀，很自然就想到了前缀树。具体来说，就是在前缀树的每个节点存储该前缀对应的值。

Python实现

# 前缀树
class TrieNode:
    def __init__(self):
        self.val = 0
        self.next = [None for _ in range(26)]


class MapSum:
    def __init__(self):
        self.root = TrieNode()
        self.map = {}


    def insert(self, key: str, val: int) -> None:
        value = val
        if key in self.map:
            value -= self.map[key]
        self.map[key] = val
        node = self.root
        for c in key:
            if node.next[ord(c) - ord('a')] is None:
                node.next[ord(c) - ord('a')] = TrieNode()
            node = node.next[ord(c) - ord('a')]
            node.val += value


    def sum(self, prefix: str) -> int:
        node = self.root
        for c in prefix:
            if node.next[ord(c) - ord('a')] is None:
                return 0
            node = node.next[ord(c) - ord('a')]
        return node.val

Java实现

class MapSum {
    class TrieNode {
        int val = 0;
        TrieNode[] next = new TrieNode[26];
    }
    
    TrieNode root;
    Map<String, Integer> map;

    public MapSum() {
        root = new TrieNode();
        map = new HashMap<>();
    }
    
    public void insert(String key, int val) {
        int value = val - map.getOrDefault(key, 0);
        map.put(key, val);
        TrieNode node = root;
        for (char c: key.toCharArray()) {
            if (node.next[c - 'a'] == null) {
                node.next[c - 'a'] = new TrieNode();
            }
            node = node.next[c - 'a'];
            node.val += value;
        }
    }
    
    public int sum(String prefix) {
        TrieNode node = root;
        for (char c: prefix.toCharArray()) {
            if (node.next[c - 'a'] == null) {
                return 0;
            }
            node = node.next[c - 'a'];
        }
        return node.val;
    }
}