В этом уроке ты
- Решишь интересную задачу из собеседования в крупной IT-компании.
- Ты поймёшь, как порой первая мысль с привычным паттерном может ввести в заблуждение и привести к не оптимальному решению.
Ты на собеседовании
Представь, что ты на собеседовании и у тебя есть 20 минут, чтобы придумать решение задачи!
Дан массив nums. Нужно найти индекс такого элемента, что сумма всех элементов слева от него равна сумме всех элементов справа.
Особые случаи:
- Если индекс указывает на первый элемент, сумма слева считается равной
0. - Если индекс указывает на последний элемент, сумма справа равна
0. - Если такого индекса нет, вернуть
1.
Пример 1
Ввод: nums = [7, -1, 4, 3, 5, 5] Вывод: 3 Объяснение: 7 + (-1) + 4 = 5 + 5
Пример 2
Ввод: nums = [0, 0, 0] Вывод: 0 Объяснение: Ответом может быть 0, 1 или 2, но возвращаем наименьший индекс.
В чем сложность задачи?
Одно из очевидных решений — использовать префиксные и суффиксные суммы.
- Префиксный массив: каждый элемент равен сумме всех предыдущих элементов исходного массива.
- Суффиксный массив: каждый элемент равен сумме всех последующих элементов исходного массива.
Например, если nums = [7, -1, 4, 3, 5, 5]
- Префиксный массив
px = [0, 7, 6, 10, 13, 18, 23] - Суффиксный массив
sx = [23, 16, 17, 13, 10, 5, 0]
Чтобы найти нужный индекс, сравниваем элементы: px[i + 1] == sx[i].
using System;
using System.Collections.Generic;
public class Solution
{
// Время: O(n), где n - длина массива nums
// Память: O(n)
public static int PivotIndex(List<int> nums)
{
// Строим префиксный массив
int[] px = new int[nums.Count + 1];
for (int i = 0; i < nums.Count; i++)
{
px[i + 1] = px[i] + nums[i];
}
// Строим суффиксный массив
int[] sx = new int[nums.Count + 1];
for (int i = nums.Count - 1; i >= 0; i--)
{
sx[i] = sx[i + 1] + nums[i];
}
// Ищем ответ
for (int i = 0; i < nums.Count; i++)
{
if (px[i + 1] == sx[i])
{
return i;
}
}
return -1; // Если не найден индекс
}
}#include <vector>
using namespace std;
// Время: O(n), где n - длина массива nums
// Память: O(n)
int pivotIndex(vector<int>& nums) {
// Строим префиксный массив
vector<int> px(nums.size() + 1, 0);
for (int i = 0; i < nums.size(); ++i) {
px[i + 1] = px[i] + nums[i];
}
// Строим суффиксный массив
vector<int> sx(nums.size() + 1, 0);
for (int i = nums.size() - 1; i >= 0; --i) {
sx[i] = sx[i + 1] + nums[i];
}
// Ищем ответ
for (int i = 0; i < nums.size(); ++i) {
if (px[i + 1] == sx[i]) {
return i;
}
}
return -1; // Если не найден индекс
}package main
import "fmt"
// Время: O(n), где n - длина массива nums
// Память: O(n)
func pivotIndex(nums []int) int {
// Строим префиксный массив
px := make([]int, len(nums)+1)
for i, num := range nums {
px[i+1] = px[i] + num
}
// Строим суффиксный массив
sx := make([]int, len(nums)+1)
for i := len(nums) - 1; i >= 0; i-- {
sx[i] = sx[i+1] + nums[i]
}
// Ищем ответ
for i := 0; i < len(nums); i++ {
if px[i+1] == sx[i] {
return i
}
}
return -1 // Если не найден индекс
}import java.util.Arrays;
public class Solution {
// Время: O(n), где n - длина массива nums
// Память: O(n)
public static int pivotIndex(int[] nums) {
// Строим префиксный массив
int[] px = new int[nums.length + 1];
for (int i = 0; i < nums.length; i++) {
px[i + 1] = px[i] + nums[i];
}
// Строим суффиксный массив
int[] sx = new int[nums.length + 1];
for (int i = nums.length - 1; i >= 0; i--) {
sx[i] = sx[i + 1] + nums[i];
}
// Ищем ответ
for (int i = 0; i < nums.length; i++) {
if (px[i + 1] == sx[i]) {
return i;
}
}
return -1; // Если не найден индекс
}
}from typing import List
# Время: O(n), где n - длина массива nums
# Память: O(n)
def pivot_index(nums: List[int]) -> int:
# Строим префиксный массив
px = [0]
for num in nums:
px.append(px[-1] + num)
# Строим суффиксный массив
sx = [0] * (len(nums) + 1)
for i in reversed(range(len(nums))):
sx[i] = sx[i + 1] + nums[i]
# Ищем ответ
for i in range(len(nums)):
if px[i + 1] == sx[i]:
return i
return -1// Время: O(n), где n - длина массива nums
// Память: O(n)
function pivotIndex(nums) {
// Строим префиксный массив
const px = [0];
for (let num of nums) {
px.push(px[px.length - 1] + num);
}
// Строим суффиксный массив
const sx = new Array(nums.length + 1).fill(0);
for (let i = nums.length - 1; i >= 0; i--) {
sx[i] = sx[i + 1] + nums[i];
}
// Ищем ответ
for (let i = 0; i < nums.length; i++) {
if (px[i + 1] === sx[i]) {
return i;
}
}
return -1; // Если не найден индекс
}Недостаток такого подхода, что он использует дополнительную память для хранения массивов префиксных и суффиксных сумм, что не является оптимальным!
Оптимальное решение
Задачу можно решить за O(n) времени и O(1) памяти.
using System;
using System.Collections.Generic;
public class Solution
{
// Время: O(n)
// Память: O(1)
public static int PivotIndex(List<int> nums)
{
int totalSum = 0;
foreach (int num in nums)
{
totalSum += num;
}
int leftSum = 0;
for (int i = 0; i < nums.Count; i++)
{
int rightSum = totalSum - leftSum - nums[i];
if (leftSum == rightSum)
{
return i;
}
leftSum += nums[i];
}
return -1; // Если не найден индекс
}
}#include <vector>
using namespace std;
// Время: O(n)
// Память: O(1)
int pivotIndex(vector<int>& nums) {
int total_sum = 0;
for (int num : nums) {
total_sum += num;
}
int left_sum = 0;
for (size_t i = 0; i < nums.size(); ++i) {
int right_sum = total_sum - left_sum - nums[i];
if (left_sum == right_sum) {
return i;
}
left_sum += nums[i];
}
return -1;
}package main
import "fmt"
// Время: O(n)
// Память: O(1)
func pivotIndex(nums []int) int {
totalSum := 0
for _, num := range nums {
totalSum += num
}
leftSum := 0
for i, num := range nums {
rightSum := totalSum - leftSum - num
if leftSum == rightSum {
return i
}
leftSum += num
}
return -1
}import java.util.List;
public class Solution {
// Время: O(n)
// Память: O(1)
public static int pivotIndex(List<Integer> nums) {
int totalSum = 0;
for (int num : nums) {
totalSum += num;
}
int leftSum = 0;
for (int i = 0; i < nums.size(); i++) {
int num = nums.get(i);
int rightSum = totalSum - leftSum - num;
if (leftSum == rightSum) {
return i;
}
leftSum += num;
}
return -1;
}
}from typing import List
# Время: O(n)
# Память: O(1)
def pivot_index(nums: List[int]) -> int:
total_sum = sum(nums)
left_sum = 0
for i, num in enumerate(nums):
right_sum = total_sum - left_sum - num
if left_sum == right_sum:
return i
left_sum += num
return -1// Время: O(n)
// Память: O(1)
function pivotIndex(nums) {
const totalSum = nums.reduce((acc, num) => acc + num, 0);
let leftSum = 0;
for (let i = 0; i < nums.length; i++) {
const num = nums[i];
const rightSum = totalSum - leftSum - num;
if (leftSum === rightSum) {
return i;
}
leftSum += num;
}
return -1;
}Идея:
- Сначала вычисляем общую сумму элементов массива:
total_sum. - Затем, проходя по массиву, накапливаем сумму элементов слева:
left_sum. - Сумму справа вычисляем как:
right_sum=total_sum−left_sum−nums[i]. - Если в какой-то момент выполняется условие
left_sum==right_sumвозвращаем текущий индекс.
Этот подход избавляет от необходимости создавать дополнительные массивы и работает оптимально. Именно такое решение чаще всего ожидает интервьюер.
Что дальше
Задачи на суммы не ограничиваются одномерными массивами. В следующем уроке ты узнаешь, как применять эту технику к более сложным структурам данных и задачам, чтобы уверенно чувствовать себя даже на самых сложных собеседованиях!