Update Custom Sort String.java

Author: varunu28

Medium/Custom Sort String.java

@@ -1,26 +1,16 @@
 class Solution {
-  public String customSortString(String S, String T) {
-    Set<Character> set = new HashSet<>();
-    for (int i = 0; i < S.length(); i++) {
-      set.add(S.charAt(i));
-    }
+  public String customSortString(String order, String s) {
     Map<Character, Integer> map = new HashMap<>();
-    StringBuilder sb = new StringBuilder();
-    for (char c : T.toCharArray()) {
-      if (set.contains(c)) {
-        map.put(c, map.getOrDefault(c, 0) + 1);
-      }
-      else {
-        sb.append(c);
-      }
-    }
-    for (int i = 0; i < S.length(); i++) {
-      char c = S.charAt(i);
-      int count = map.getOrDefault(c, 0);
-      while (count-- > 0) {
-        sb.append(c);
-      }
+    int position = 1;
+    for (char c : order.toCharArray()) {
+      map.put(c, position++);
     }
-    return sb.toString();
+    return s.chars()
+        .mapToObj(c -> (char) c)
+        .sorted(Comparator.comparing(o -> map.getOrDefault(o, 0)))
+        .collect(Collector.of(StringBuilder::new,
+          StringBuilder::append,
+          StringBuilder::append,
+          StringBuilder::toString));
   }
 }