第三次算法笔试 蚂蚁

蚂蚁春招

一共三道算法，很菜，只做出第一道最简单的

第一题忘记了

第二题是一个二叉树的题目，感觉是很简单的，但人太菜，没办法，复习一下二叉树

题目大概如下：

输入如下：

5 2  # 5个节点，问询 2次
1 2  # 节点1增加子节点 2
1 3  # 节点1增加子节点 3
2 4
2 5
3 6  # 第一次问询，节点3与节点6的距离是多少
1 5  # 第二次问询，节点1与节点5的距离是多少

第三题：给出一个数组an = [a1, a2, ..., an]

题目要求直接写的话是这样：但肯定需要优化一下算法才行

an = [1,2,5,6,2,1,5,2222,11111,2222， ...]
sum = 0
for i in an:
	for j in an:
		sum += i //j
print(sum)

优化后的代码（来自deepseek）使用_small_m_or_low_max就足够了，还是见得少了，这个没想到

from collections import defaultdict
import bisect
import math

def dynamic_optimized_sum(an):
    freq = defaultdict(int)
    for num in an:
        freq[num] += 1
    sorted_j = sorted(freq.keys())
    m = len(sorted_j)
    max_i = max(sorted_j) if m > 0 else 0
    
    # 选择策略：根据m和max_i决定算法
    if m <= 1000 or max_i <= 1e5:
        return _small_m_or_low_max(freq, sorted_j)
    else:
        return _large_m_or_high_max(freq, sorted_j, max_i)

def _small_m_or_low_max(freq, sorted_j):
    total = 0
    for j_val in sorted_j:
        if j_val == 0:
            continue
        cnt_j = freq[j_val]
        sum_i = 0
        for i_val in freq:
            sum_i += (i_val // j_val) * freq[i_val]
        total += sum_i * cnt_j
    return total

def _large_m_or_high_max(freq, sorted_j, max_i):
    prefix = [0] * (len(sorted_j) + 1)
    for i in range(len(sorted_j)):
        prefix[i + 1] = prefix[i] + freq[sorted_j[i]]
    total = 0
    for i_val in freq:
        cnt_i = freq[i_val]
        if i_val == 0:
            continue
        sum_j_part = 0
        sqrt_i = int(math.isqrt(i_val))
        # 处理小j_val (j <= sqrt(i_val))
        pos = bisect.bisect_right(sorted_j, sqrt_i)
        for j in sorted_j[:pos]:
            if j == 0:
                continue
            q = i_val // j
            sum_j_part += q * freq[j]
        # 处理大j_val (j > sqrt(i_val))，数论分块
        for q in range(1, sqrt_i + 1):
            lower = (i_val // (q + 1)) + 1
            upper = i_val // q
            if lower > upper:
                continue
            left = bisect.bisect_left(sorted_j, lower)
            right = bisect.bisect_right(sorted_j, upper)
            count = prefix[right] - prefix[left]
            sum_j_part += q * count
        total += sum_j_part * cnt_i
    return total