[Java 8+] Map과 merge(), compute()

[Java 8+] Map.merge(), Map.compute()

Map은 백엔드 개발에서도 자주 사용되는 컬렉션이고, PS 문제를 풀 때도 자주 등장하는 컬렉션이다.
키의 존재 유무를 파악하고, 그 값에 따라 저장하는 값이 달라지는 형태의 문제를 자주 봤을 것이다.

 

자바의 Map 라이브러리에는 키가 이미 존재할 때 값 갱신을 간편하게 처리할 수 있는 유용한 메서드들이 있다. 그중 대표적인 것이 merge와 compute이다.

Merge

Map.merge(K key, V value, BiFunction<? super V,? super V,? extends V> remappingFunction)

• 특정 키에 대해 새 값을 병합하는 메서드이다.
• 동작 방식은 다음과 같다.
  1. 키가 존재하지 않으면 value를 그대로 삽입한다.
  2. 키가 존재하면 기존 값과 새로운 value를 remappingFunction으로 병합하여 새로운 값을 저장한다

Map<String, Integer> map = new HashMap<>();
map.put("apple", 2);

// 키 apple이 이미 존재하므로, 2와 3을 더해 새로운 값 5를 저장한다.
map.merge("apple", 3, Integer::sum);  

// 키 banana가 없으므로, 새로운 값 1을 그대로 삽입한다.
map.merge("banana", 1, Integer::sum);

System.out.println(map); // {apple=5, banana=1}

 

Compute

Map.compute(K key, BiFunction<? super K,? super V,? extends V> remappingFunction)

• 특정 키에 대해 값을 계산하여 저장하는 메서드이다.
• 동작 방식은 다음과 같다.
  1. 키가 존재하면 remappingFunction에 (key, oldValue)를 전달한다.
  2. 키가 존재하지 않으면 remappingFunction에 (key, null)을 전달한다.
  3. 반환된 값이 null이면 해당 키를 맵에서 제거한다.

Map<String, Integer> map = new HashMap<>();
map.put("apple", 2);

// 기존 값(2)에 대해 연산을 수행하여 4로 갱신한다.
map.compute("apple", (k, v) -> v * 2);

// 키가 없는 경우 null로 들어오므로 기본값을 할당한다.
map.compute("banana", (k, v) -> v == null ? 1 : v + 1);

System.out.println(map); // {apple=4, banana=1}

 

PS 활용 예시(BOJ 7662. 이중 우선순위 큐)

이 문제 뿐만 아니라, Map을 사용하는 대부분의 문제에서 위 메서드를 사용하면 코드를 깔끔하게 작성할 수 있다.

Before

TreeMap<Integer, Integer> tree = new TreeMap<>();
int k = Integer.parseInt(br.readLine());

while (k-- > 0) {
    String[] parts = br.readLine().split(" ");
    if (parts[0].equals("I")) {
        int val = Integer.parseInt(parts[1]);
        int cur = tree.getOrDefault(val, 0);
        tree.put(val, cur + 1);
    } else if (tree.isEmpty()) {
        continue;
    } else {
        int key = 0;
        if (parts[1].equals("-1")) {
            key = tree.firstKey();
        } else {
            key = tree.lastKey();
        }
        int val = tree.get(key);
        if (val == 1) {
            tree.remove(key);
        } else {
            tree.put(key, val - 1);
        }
    }
}

After

while (t-- > 0) {
    TreeMap<Integer, Integer> tree = new TreeMap<>();
    int k = Integer.parseInt(br.readLine());

    while (k-- > 0) {
        String[] parts = br.readLine().split(" ");
        if (parts[0].equals("I")) {
            int val = Integer.parseInt(parts[1]);
            tree.merge(val, 1, Integer::sum);
        } else if (!tree.isEmpty()) {
            if (parts[1].equals("-1")) {
                tree.compute(tree.firstKey(), (k1, v) -> (v == 1) ? null : v - 1);
            } else {
                tree.compute(tree.lastKey(), (k1, v) -> (v == 1) ? null : v - 1);
            }
        }
    }
    if (tree.isEmpty()) {
        sb.append("EMPTY\n");
    } else {
        sb.append(tree.lastKey()).append(' ').append(tree.firstKey()).append('\n');
    }
}

 

코드가 매우 짧고 간결하게 정리된 것을 확인할 수 있다.