Google | Onsite | K largest values with labels
Sithis
2436
GoogleOnsiteInterview

Given an array of values and their corresponding labels. Find out the k largest values, but each label cannot be used more than m times.

Example 1:

Input: values = [9, 8, 6, 8, 7], labels = [A, A, B, A, B], k = 3, m = 2
Output: [9, 8, 7]
Explanation: 8 > 7, but we can't use the value with A label more than 2 times.
Follow-up 1

Do beter than O(nlogn).

Follow-up 2

Can you do it in O(n) time?

My Java solutions

https://leetcode.com/playground/mtwVrCuH
Heap O(nlogk):

public static List<Integer> topK(int[] values, char[] labels, int k, int m) {
	Map<Character, Queue<Integer>> map = new HashMap<>();
	for (int i = 0; i < values.length; i++) {
		char label = labels[i];
		Queue<Integer> pq = map.getOrDefault(label, new PriorityQueue<>(m + 1));
		pq.add(values[i]);
		if (pq.size() > m) pq.poll();
		map.put(label, pq);
	} 

	Queue<Integer> result = new PriorityQueue<>(k + 1);
	for (Queue<Integer> pq : map.values()) {
		for (int val : pq) {
			result.add(val);
			if (result.size() > k) result.poll();
		}   
	}
	return new ArrayList<>(result);
}

Qiuckselect O(n):

public static List<Integer> topK2(int[] values, char[] labels, int k, int m) {
	Map<Character, List<Integer>> map = new HashMap<>();
	for (int i = 0; i < values.length; i++) {
		char label = labels[i];
		map.putIfAbsent(label, new ArrayList<>());
		map.get(label).add(values[i]);
	}

	List<Integer> result = new ArrayList<>();
	for (List<Integer> vals : map.values()) {
		qselect(vals, m);
		result.addAll(vals.subList(0, Math.min(m, vals.size())));
	}

	qselect(result, k);
	return result.subList(0, Math.min(k, result.size()));
}

private static void qselect(List<Integer> arr, int k) {
	if (k >= arr.size()) return;
	k--; // k zero baised for convenience
	int lo = 0;
	int hi = arr.size() - 1;
	while (lo < hi) {
		int p = partition(arr, lo, hi);
		if (p < k) {
			lo = p + 1;
		} else if (p > k) {
			hi = p - 1;
		} else {
			break;
		}
	}
}

private static int partition(List<Integer> arr, int lo, int hi) {
	int i = lo;
	int j = hi + 1;
	int pivot = arr.get(lo);
	while (true) {
		while (i < hi && arr.get(++i) > pivot);
		while (arr.get(--j) < pivot);

		if (i >= j) break;
		Collections.swap(arr, i, j);
	}
	Collections.swap(arr, lo, j);
	return j;
}

https://leetcode.com/problems/largest-values-from-labels

Comments (11)