배열 인덱스가 노래의 고유 번호이고, 장르 목록과 플레이 수 목록이 주어졌을 때, 장르별로 가장 많이 재생된 노래를 두 개 씩 선택하는 문제이다.

  1. 속한 노래가 많이 재생된 장르를 먼저 수록
  2. 장르 내에서 많이 재생된 노래를 먼저 수록
  3. 장르 내에서 재생 횟수가 같은 노래 중에서는 고유 번호가 낮은 노래를 먼저 수록

장르에 속한 곡이 하나라면, 하나의 곡만 선택한다. 노래의 개수는 10000이하이므로 O(N^2) 내에서 해결한다.

public int[] solution(String[] genres, int[] plays) {
    Map<String, Integer> sumMap = new HashMap<>();

    for (int i = 0; i < genres.length; i++) {
        sumMap.put(genres[i], sumMap.getOrDefault(genres[i], 0) + plays[i]);
    }
    
    Map<String, List<Integer>> indexMap = IntStream.range(0, plays.length)
        .boxed()
        .sorted((o1, o2) -> plays[o2] - plays[o1])
        .collect(Collectors.groupingBy(integer -> genres[integer]));

    return indexMap
        .keySet()
        .stream()
        .sorted(Comparator.comparingInt(sumMap::get).reversed())
        .reduce(new ArrayList<>(), (List<Integer> integers, String s) -> {
            List<Integer> list = indexMap.get(s);
            integers.add(list.get(0));
            if (list.size() > 1) {
                integers.add(list.get(1));
            }
            return integers;
        }, (integers, integers2) -> {
            integers.addAll(integers2);
            return integers;
        })
        .stream()
        .mapToInt(value -> value)
        .toArray();
}

이번엔 자바8 stream을 사용해서 풀어봤다. 정렬하면서 요소 몇개를 선택하고 이런 연속적인 기능을 표현하는데 효과적인 것 같다. 위의 코드가 초기 코드였는데, 테스트는 통과했지만 만족스럽지가 않다.

 private int[] bestAlbum(String[] genres, int[] plays) {
     return IntStream.range(0, plays.length)
         .boxed()
         .collect(Collectors.groupingBy(id -> genres[id]))
         .entrySet()
         .stream()
         .sorted((e1, e2) -> sum(plays, e2.getValue()) - sum(plays, e1.getValue()))
         .flatMap(e -> e.getValue().stream().sorted((id1, id2) -> {
             if (plays[id2] == plays[id1]) return id1 - id2;
             return plays[id2] - plays[id1];
         }).limit(2))
         .mapToInt(value -> value)
         .toArray();
 }

private int sum(int[] plays, List<Integer> ids) {
    return ids.stream()
        .mapToInt(i -> plays[i])
        .sum();
}

IntStream.range를 사용해서 필요없는 Map을 없애고 keySet이 아닌 entrySet와 flatMap을 사용하여 복잡해보이는 컬렉션 생성 과정을 없앴다.

private int[] solution(String[] genres, int[] plays) {
    return IntStream.range(0, genres.length)
        .mapToObj(i -> new Music(i, genres[i], plays[i]))
        .collect(Collectors.groupingBy(Music::getGenre))
        .entrySet()
        .stream()
        .sorted((o1, o2) -> sum(o2.getValue()) - sum(o1.getValue()))
        .flatMap(e -> e.getValue().stream().sorted().limit(2))
        .mapToInt(m -> m.id).toArray();
}

private int sum(List<Music> value) {
    return value
        .stream()
        .reduce(0, (sum, music) -> sum + music.played, (sum1, sum2) -> sum1 + sum2);
}

class Music implements Comparable<Music> {
    int id;
    String genre;
    int played;

    public Music(int id, String genre, int play) {
        this.id = id;
        this.genre = genre;
        this.played = play;
    }

    @Override
    public int compareTo(Music o) {
        if(played == o.played) return id - o.id;
        return o.played - played;
    }

    public String getGenre() {
        return genre;
    }
}

비슷하게 stream을 사용하지만, 알아보기 힘든 배열 접근 코드를 없애고 객체와 Comparable 인터페이스를 활용하는 방법이다.