logoChat01.ai

как можно решить такую задачу за O(nlogn) или мень...

創建於：2025年5月9日

使用 O3 回答 Chat01

問題

как можно решить такую задачу за O(nlogn) или меньше:
У мальчика Олега сегодня День Рождения, ему исполняется 10 лет. Так как 10 лет особый возраст, то родители подарили ему особый подарок: массив из чисел от 1 до 10. Обрадовавшись такому подарку, он сразу начал изучать его свойства.
Недавно на прогулке друзья рассказали Олегу о арифметических прогрессиях. Для изучения свойств массива Олег решил начать с чего-то малого, а именно с арифметических прогрессий длиной 3. Три числа a, b, с образуют арифметическую прогрессию, если b - a = c - b. Оказалось, что не все подогрезки массива могут содержать такие арифметические прогрессии, поэтому Олег решил вначале узнать, сколько существует подогрезков, в которых можно выбрать три числа, не меняя их порядок, с нужным свойством.
Подотрезком массива A от l до r называется массив из элементов A_l, A_(l+1), ..., A_r.
Пример
Ввод:
5
1 2 3 3 2
Вывод:
3
Формат входных данных
Дано число n - длина массива Олега (3 <= n <= 10^5)
И даны n целых чисел - подаренный массив (1 <= A_i <= 10)
НО: учитывай то, что в задаче требуется найти все подотрезки, в которых есть хотя бы одна такая тройка чисел, то есть если подотрезок содержит эту тройку, то все подотрезки, содержащие этот подотрезок - тоже будут учитываться в счетчике
например:
массив 2 1 2 3 7 10 17
возьмем подотрезок 1 2 3 - это нужная нам тройка, значит подотрезки 2 1 2 3, 2 1 2 3 7, 1 2 3 7, 2 1 2 3 7 10, 1 2 3 7 10, … тоже должны учитываться

Python

python
def count_subarrays_with_ap(arr):
    n=len(arr)
    V=11
    last_before=[[-1]*V for _ in range(n)]
    last=[-1]*V
    for i in range(n):
        val=arr[i]
        last[val]=i
        for v in range(1,V):
            last_before[i][v]=last[v]
    next_after=[[-1]*V for _ in range(n)]
    nxt=[-1]*V
    for i in range(n-1,-1,-1):
        val=arr[i]
        nxt[val]=i
        for v in range(1,V):
            next_after[i][v]=nxt[v]
    min_r=[None]*n
    for j in range(n):
        b=arr[j]
        for d in range(1,10):
            a=b-d
            c=b+d
            if a>=1 and c<=10:
                i= last_before[j-1][a] if j-1>=0 else -1
                k= next_after[j+1][c] if j+1<n else -1
                if i!=-1 and k!=-1:
                    if min_r[i] is None or k<min_r[i]:
                        min_r[i]=k
    # propagate backwards
    for l in range(n-2,-1,-1):
        if min_r[l] is None or (min_r[l+1] is not None and min_r[l+1]<min_r[l]):
            min_r[l]=min_r[l+1]
    total=0
    for l in range(n):
        if min_r[l] is not None:
            total += n - min_r[l]
    return total

print(count_subarrays_with_ap([1,2,3,3,2]))

Python

python
def count_subarrays_with_ap(arr):
    n=len(arr)
    V=11
    last_before=[[-1]*V for _ in range(n)]
    last=[-1]*V
    for i in range(n):
        last[arr[i]]=i
        for v in range(1,V):
            last_before[i][v]=last[v]
    next_after=[[-1]*V for _ in range(n)]
    nxt=[-1]*V
    for i in range(n-1,-1,-1):
        nxt[arr[i]]=i
        for v in range(1,V):
            next_after[i][v]=nxt[v]
    min_r=[None]*n
    for j in range(n):
        b=arr[j]
        # diff 0
        i = last_before[j-1][b] if j-1>=0 else -1
        k = next_after[j+1][b] if j+1<n else -1
        if i!=-1 and k!=-1:
            if min_r[i] is None or k<min_r[i]:
                min_r[i]=k
        # diff 1..9
        for d in range(1,10):
            a=b-d
            c=b+d
            if 1<=a<=10 and 1<=c<=10:
                i= last_before[j-1][a] if j-1>=0 else -1
                k= next_after[j+1][c] if j+1<n else -1
                if i!=-1 and k!=-1:
                    if min_r[i] is None or k<min_r[i]:
                        min_r[i]=k
    for l in range(n-2,-1,-1):
        if min_r[l] is None or (min_r[l+1] is not None and min_r[l+1]<min_r[l]):
            min_r[l]=min_r[l+1]
    total=0
    for l in range(n):
        if min_r[l] is not None:
            total += n - min_r[l]
    return total
print(count_subarrays_with_ap([1,1,1]))

Python

python
import random, itertools
def brute(arr):
    n=len(arr)
    cnt=0
    for l in range(n):
        for r in range(l,n):
            ok=False
            sub=arr[l:r+1]
            m=r-l+1
            if m>=3:
                # equivalent to subsequence choose 3
                for i in range(m):
                    for j in range(i+1,m):
                        for k in range(j+1,m):
                            if sub[j]-sub[i]==sub[k]-sub[j]:
                                ok=True
                                break
                        if ok: break
                    if ok: break
            if ok:
                cnt+=1
    return cnt
for _ in range(100):
    n=random.randint(3,8)
    arr=[random.randint(1,10) for _ in range(n)]
    if brute(arr)!=count_subarrays_with_ap(arr):
        print(arr, brute(arr), count_subarrays_with_ap(arr))
        break

Python

python
def find_triples(arr):
    n=len(arr)
    triples=[]
    for i in range(n):
        for j in range(i+1,n):
            for k in range(j+1,n):
                if arr[j]-arr[i]==arr[k]-arr[j]:
                    triples.append((i,j,k,arr[i],arr[j],arr[k]))
    return triples
triples=find_triples([7, 10, 8, 1, 7, 4])
triples

[(1, 4, 5, 10, 7, 4)]

Python

python
def count_subarrays_with_ap2(arr):
    n=len(arr)
    V=11
    last_before=[[-1]*V for _ in range(n)]
    last=[-1]*V
    for i in range(n):
        last[arr[i]]=i
        for v in range(1,V):
            last_before[i][v]=last[v]
    next_after=[[-1]*V for _ in range(n)]
    nxt=[-1]*V
    for i in range(n-1,-1,-1):
        nxt[arr[i]]=i
        for v in range(1,V):
            next_after[i][v]=nxt[v]
    min_r=[None]*n
    for j in range(n):
        b=arr[j]
        for a in range(1,11):
            if a==b:
                # need at least one previous occurrence
                i= last_before[j-1][a] if j-1>=0 else -1
            else:
                i= last_before[j-1][a] if j-1>=0 else -1
            if i==-1: continue
            d=b - a
            c=b + d
            if 1<=c<=10:
                k= next_after[j+1][c] if j+1<n else -1
                if k!=-1:
                    if min_r[i] is None or k<min_r[i]:
                        min_r[i]=k
    for l in range(n-2,-1,-1):
        if min_r[l] is None or (min_r[l+1] is not None and min_r[l+1]<min_r[l]):
            min_r[l]=min_r[l+1]
    total=0
    for l in range(n):
        if min_r[l] is not None:
            total += n - min_r[l]
    return total

print(count_subarrays_with_ap2([7,10,8,1,7,4]))

思考

def count_subarrays_with_ap_bruteprecomputed(arr): n=len(arr) total=0 for l in range(n): for r in range(l,n): ok=False if r-l+1>=3: for i in range(l,r-1): for j in range(i+1,r): for k in range(j+1,r+1): if arr[j]-arr[i]==arr[k]-arr[j]: ok=True break if ok: break if ok: break if ok: total+=1 return total

for _ in range(100): n=random.randint(3,10) arr=[random.randint(1,10) for _ in range(n)] if count_subarrays_with_ap2(arr)!=count_subarrays_with_ap_bruteprecomputed(arr): print("Mismatch", arr, count_subarrays_with_ap2(arr), count_subarrays_with_ap_bruteprecomputed(arr)) break

Thought for 3m 38s

回答

Идея — считать «плохие» подотрезки сложно, зато легко «приклеивать» все минимальные отрезки, где тройка уже нашлась

Пусть

i < j < k,        A[j]-A[i] = A[k]-A[j]

— индексы одной подходящей тройки.
Тогда каждый подотрезок, чей левый конец ≤ i, а правый ≥ k, автоматически «хороший».
Если для каждого i мы будем знать самый левый возможный k, то для фиксированного i число подходящих правых концов равно n – k + 1.

Задача сводится к тому, чтобы для каждого индекса i найти

text
minR[i] = минимальный k ≥ i,
          для которого существует j (i < j < k),
          образующий с i и k арифметическую прогрессию

а затем сосчитать
Σ_{i=1..n} (n – minR[i] + 1).

Как найти все minR[i] за O(n·σ), где σ = 10

  1. Предподсчёты

    • Значения лежат в диапазоне 1…10 ⇒ удобно хранить позиции для каждого значения.
    • Построим два массива 
      text
      last[v][p]  – индекс последнего v слева от позиции p
next[v][p]  – индекс ближайшего v справа от позиции p
      Их можно заполнить одним левым и одним правым проходом, затратив O(n·σ) памяти и времени.

  2. Перебор середины j
    Для каждой позиции j (значение b = A[j]) перебираем все возможные значения слева a (всего 10).

    • Если слева от j есть a, берём его последний индекс i = last[a][j-1].
    • Тогда шаг прогрессии d = b − a, а кандидат справа — c = b + d = 2b − a.
      Если c попадает в диапазон 1…10 и справа от j существует такой c, берём его первый индекс k = next[c][j+1].
    • Пара (i,k) образует минимальный подотрезок, начинающийся в i, где уже найдена тройка.
      Сохраняем minR[i] = min(minR[i], k).

    При σ = 10 на каждую позицию выполняем ≤ 10 операций ⇒ O(n·σ).

    Разница шага может быть отрицательной → арифметическая прогрессия может убывать.
    Перебирая все a мы автоматически учитываем и положительные, и отрицательные d,
    а случай d = 0 (три одинаковых числа) обрабатывается тем же кодом, если слева и справа от j есть такое же значение.

  3. Распространение вправо
    minR[i] может быть None, если тройка начинается позже.
    Проходим массив справа-налево:

    text
    for i = n-2 … 0:
     if minR[i] is None or (minR[i+1] < minR[i]):
             minR[i] = minR[i+1]

    Теперь minR[i] - это наименьший правый конец для любого подотрезка, начинающегося не левее i.

  4. Подсчёт ответа

    ans = Σ_{i=0}^{n-1}  ( n - minR[i] )

    (слагаем только те i, где minR[i] определено).

Ассимптотики

  • Подготовка last и next — O(n·σ)
  • Перебор «центров» — O(n·σ)
  • Один линейный проход для распространения — O(n)
  • Память — O(n·σ)

При σ = 10 это O(n) по времени и ≈ 10 · n по памяти, что удовлетворяет требованию O(n log n).

Краткая иллюстрация (пример из условия)

i12345
A12332
  • Тройка 1 2 3 даёт пару (i=1, k=3) → minR[1] = 3.
  • После «проталкивания» влево получаем
    minR = [3,3,-,-,-].
  • Вклад:
    • для i=1 → 5-3+1 = 3 подотрезка
    • для i=0 → тот же 3
  • Ответ = 3, как в примере.

Простой код на C++ 17

cpp
#include <bits/stdc++.h>
using namespace std;
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n;  if (!(cin >> n)) return 0;
    vector<int> a(n);
    for (int &x: a) cin >> x;

    const int V = 11;                 // значения 1‥10
    vector<array<int,V>> last(n), nxt(n);
    array<int,V> cur;  cur.fill(-1);

    // last[v][i] – последний v слева от i
    for (int i = 0; i < n; ++i) {
        cur[a[i]] = i;
        last[i]    = cur;
    }
    // nxt[v][i] – первый v справа от i
    cur.fill(-1);
    for (int i = n-1; i >= 0; --i) {
        cur[a[i]] = i;
        nxt[i]    = cur;
    }

    const long long INF = 1e18;
    vector<int> minR(n, -1);          // минимальный k для данного i

    for (int j = 0; j < n; ++j) {
        int b = a[j];
        // перебираем ВСЕ возможные a (достаточно 1‥10)
        for (int av = 1; av < V; ++av) {
            int i = (j ? last[j-1][av] : -1);
            if (i == -1) continue;
            int d = b - av;
            int cv = b + d;                 // 2b - a
            if (cv < 1 || cv > 10) continue;
            int k = (j+1 < n ? nxt[j+1][cv] : -1);
            if (k == -1) continue;
            if (minR[i] == -1 || k < minR[i]) minR[i] = k;
        }
    }

    // распространяем минимумы вправо
    for (int i = n-2; i >= 0; --i) {
        if (minR[i] == -1 || (minR[i+1] != -1 && minR[i+1] < minR[i]))
            minR[i] = minR[i+1];
    }

    long long ans = 0;
    for (int i = 0; i < n; ++i)
        if (minR[i] != -1) ans += n - minR[i];

    cout << ans << '\n';
    return 0;
}

Алгоритм линейный, память ≈ 10 · n целых — без проблем укладывается в ограничения n ≤ 10^5.

分享這個問答