Ammar-utiful Array

Pista 1

El problema menciona prefijos de arreglos; eso ya debería resultarte familiar.

Pista 2

Si el prefijo de longitud $10$ de un color no es Ammar-utiful, ¿un prefijo de $11$ alguna vez será Ammar-utiful?

Solución

Explicación

Primero computamos sumas de prefijos para cada color para poder hallar en $\mathcal{O}(1)$ cuáles son todos los prefijos de cada color.

Las consultas de tipo $1$ son un poco más difíciles de manejar porque suman una cantidad a cada elemento que no es de cierto color. Por eso, llevar simplemente cuánto se le sumó a cada color sería demasiado lento.

Para sortearlo, en su lugar llevamos cuánto no se le sumó a cada color, junto con un contador de la cantidad total sumada.

Por ejemplo, digamos que tenemos los colores $1$ , $2$ y $3$ , junto con las siguientes consultas:


1 1 2
1 2 3

Nuestra cantidad total sumada después de esto sería $5$ . $2$ no se le habría sumado al color $1$ , y $3$ no se le habría sumado al color $2$ .

Para obtener cuánto se le sumó a un color específico, tomamos el contador total y restamos cuánto no se le sumó a ese color.

Después de todo esto, podemos hacer búsqueda binaria sobre el prefijo máximo para cada consulta de tipo $2$ .

Implementación

Complejidad temporal: $\mathcal{O}(N)$ de preprocesamiento y $\mathcal{O}(\log N)$ por cada consulta

C++Java


#include <iostream>
#include <map>
#include <vector>
 
using std::cout;
using std::endl;
using std::vector;
 
int main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(NULL);
 
	int n;
	std::cin >> n;
	vector<int> arr(n);
	for (int &i : arr) { std::cin >> i; }
	vector<int> color(n);
	for (int &c : color) { std::cin >> c; }
 
	std::map<int, vector<long long>> col_prefs;
	for (int i = 0; i < n; i++) {
		int c = color[i];
		if (!col_prefs.count(c)) { col_prefs[c] = {0}; }
		col_prefs[c].push_back(col_prefs[c].back() + arr[i]);
	}
 
	long long tot_inc = 0;
	std::map<int, long long> col_exclude;
 
	int query_num;
	std::cin >> query_num;
	for (int q = 0; q < query_num; q++) {
		int q_type, col;
		long long arg;
		std::cin >> q_type >> col >> arg;
 
		if (q_type == 1) {
			tot_inc += arg;
			col_exclude[col] += arg;
		} else if (q_type == 2) {
			int lo = 0;
			int hi = col_prefs[col].size() - 1;
			int valid = -1;
			while (lo <= hi) {
				int mid = (lo + hi) / 2;
				long long init_val = col_prefs[col][mid];
				long long to_add = (tot_inc - col_exclude[col]) * mid;
				if (init_val + to_add <= arg) {
					valid = mid;
					lo = mid + 1;
				} else {
					hi = mid - 1;
				}
			}
 
			cout << valid << '\n';
		}
	}
}


import java.io.*;
import java.util.*;
 
public class AmmarutilArray {
	public static void main(String[] args) throws IOException {
		BufferedReader read = new BufferedReader(new InputStreamReader(System.in));
 
		int n = Integer.parseInt(read.readLine());
		StringTokenizer st = new StringTokenizer(read.readLine());
		int[] arr = new int[n];
		for (int i = 0; i < n; i++) { arr[i] = Integer.parseInt(st.nextToken()); }
		st = new StringTokenizer(read.readLine());
		int[] color = new int[n];
		for (int i = 0; i < n; i++) { color[i] = Integer.parseInt(st.nextToken()); }
 
		Map<Integer, List<Long>> colPrefs = new HashMap<>();
		for (int i = 0; i < n; i++) {
			int c = color[i];
			if (!colPrefs.containsKey(c)) {
				colPrefs.put(c, new ArrayList<>(Arrays.asList(0L)));
			}
			List<Long> pref = colPrefs.get(c);
			pref.add(pref.get(pref.size() - 1) + arr[i]);
		}
 
		long totInc = 0;
		Map<Integer, Long> colExclude = new HashMap<>();
 
		StringBuilder ans = new StringBuilder();
		int queryNum = Integer.parseInt(read.readLine());
		for (int q = 0; q < queryNum; q++) {
			StringTokenizer query = new StringTokenizer(read.readLine());
			int qType = Integer.parseInt(query.nextToken());
			int col = Integer.parseInt(query.nextToken());
			long arg = Long.parseLong(query.nextToken());
 
			if (qType == 1) {
				totInc += arg;
				colExclude.put(col, colExclude.getOrDefault(col, 0L) + arg);
			} else if (qType == 2) {
				int lo = 0;
				int hi = colPrefs.get(col).size() - 1;
				int valid = -1;
				while (lo <= hi) {
					int mid = (lo + hi) / 2;
					long initVal = colPrefs.get(col).get(mid);
					long toAdd = (totInc - colExclude.getOrDefault(col, 0L)) * mid;
					if (initVal + toAdd <= arg) {
						valid = mid;
						lo = mid + 1;
					} else {
						hi = mid - 1;
					}
				}
 
				ans.append(valid).append('\n');
			}
		}
 
		System.out.print(ans);
	}
}