Top K问题系列之三 手写代码

本文深入探讨了TopK问题的不同解法及其应用场景,重点讲解了如何使用堆排序和快速排序的变种来解决这类问题。文章提供了详细的代码示例,帮助读者理解如何找到最小或最大的K个元素。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

Top K问题是面试时手写代码的常考题,某些场景下的解法与堆排快排的关系紧密,所以把它放在堆排后面讲。


关于Top K问题最全的分类总结是在这里(包括海量数据的处理),个人将这些题分成了两类:一类是容易写代码实现的;另一类侧重考察思路的。毫无疑问,后一种比较简单,你只要记住它的应用场景、解决思路,并能在面试的过程中将它顺利地表达出来,便能以不变应万变。前一种,需要手写代码,就必须要掌握一定的技巧,常见的解法有两种,就是前面说过的堆排和快排的变形。

堆排解法

用堆排来解决Top K的思路很直接。

前面已经说过,堆排利用的大(小)顶堆所有子节点元素都比父节点小(大)的性质来实现的,这里故技重施:既然一个大顶堆的顶是最大的元素,那我们要找最小的K个元素,是不是可以先建立一个包含K个元素的堆,然后遍历集合,如果集合的元素比堆顶元素小(说明它目前应该在K个最小之列),那就用该元素来替换堆顶元素,同时维护该堆的性质,那在遍历结束的时候,堆中包含的K个元素是不是就是我们要找的最小的K个元素?

实现: 
在堆排的基础上,稍作了修改,buildHeap和heapify函数都是一样的实现,不难理解。

速记口诀:最小的K个用最大堆,最大的K个用最小堆。


public class TopK {

public static void main(String[] args) {

// TODO Auto-generated method stub

int[] a = { 1, 17, 3, 4, 5, 6, 7, 16, 9, 10, 11, 12, 13, 14, 15, 8 };

int[] b = topK(a, 4);

for (int i = 0; i < b.length; i++) {

System.out.print(b[i] + ", ");

}

}

public static void heapify(int[] array, int index, int length) {

int left = index * 2 + 1;

int right = index * 2 + 2;

int largest = index;

if (left < length && array[left] > array[index]) {

largest = left;

}

if (right < length && array[right] > array[largest]) {

largest = right;

}

if (index != largest) {

swap(array, largest, index);

heapify(array, largest, length);

}

}

public static void swap(int[] array, int a, int b) {

int temp = array[a];

array[a] = array[b];

array[b] = temp;

}

public static void buildHeap(int[] array) {

int length = array.length;

for (int i = length / 2 - 1; i >= 0; i--) {

heapify(array, i, length);

}

}

public static void setTop(int[] array, int top) {

array[0] = top;

heapify(array, 0, array.length);

}

public static int[] topK(int[] array, int k) {

int[] top = new int[k];

for (int i = 0; i < k; i++) {

top[i] = array[i];

}

//先建堆,然后依次比较剩余元素与堆顶元素的大小,比堆顶小的, 说明它应该在堆中出现,则用它来替换掉堆顶元素,然后沉降。 buildHeap(top); for (int j = k; j < array.length; j++) { int temp = top[0]; if (array[j] < temp) { setTop(top, array[j]); } } return top; }}

### 手写实现注意力机制 在深度学习中,自注意力机制是一种用于捕捉序列内部依赖关系的方法。为了更好地理解并手动编写这一过程,可以将其分解成几个核心部分:计算查询(query)、键(key)和值(value),以及应用缩放点积注意函数。 #### 计算 Query、Key 和 Value 向量 对于输入的数据 X,可以通过线性变换来获得 Q、K 和 V: ```python import torch import torch.nn as nn class AttentionMechanism(nn.Module): def __init__(self, d_model, num_heads): super(AttentionMechanism, self).__init__() assert d_model % num_heads == 0 # 定义权重矩阵 W_Q, W_K, W_V self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.num_heads = num_heads self.d_k = d_model // num_heads def forward(self, q, k, v): batch_size = q.size(0) # 对 QKV 进行线性变换 q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k = self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v = self.W_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) return q, k, v ``` 不同头部携带独立的 \(W^Q\)、\(W^K\) 和 \(W^V\) 权重矩阵来进行特征转换[^2]。 #### 缩放点积注意力 一旦获得了 Q、K 和 V 的表示形式,则可通过如下方式计算注意力分数: \[ \text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \] 其中 \(d_k\) 是维度大小,用来防止梯度消失或爆炸问题。 ```python def scaled_dot_product_attention(q, k, v, mask=None): """Compute 'Scaled Dot Product Attention'""" dk = torch.tensor([k.shape[-1]], dtype=torch.float32).sqrt() scores = torch.matmul(q, k.transpose(-2, -1)) / dk if mask is not None: scores += (mask * -1e9) attention_weights = torch.softmax(scores, dim=-1) output = torch.matmul(attention_weights, v) return output, attention_weights ``` 此方法允许模型聚焦于最相关的词项上,从而增强对上下文的理解能力。 #### 多头注意力融合 最后一步是将多个平行运行的关注力子空间的结果组合起来形成最终输出向量。这通常涉及到另一个投影操作以匹配原始输入尺寸。 ```python class MultiHeadedAttention(nn.Module): def __init__(self, h, d_model, dropout=0.1): "Take in model size and number of heads." super(MultiHeadedAttention, self).__init__() assert d_model % h == 0 self.attention_mechanism = AttentionMechanism(d_model=d_model, num_heads=h) self.fc_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, q, k, v, mask=None): q, k, v = self.attention_mechanism(q, k, v) x, _ = scaled_dot_product_attention(q, k, v, mask) x = x.transpose(1, 2).contiguous().view(x.size(0), -1, self.h * self.d_k) return self.fc_o(self.dropout(x)) ``` 上述代码展示了如何构建一个多头自我关注模块,在自然语言处理任务中有广泛应用价值。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值