牛客NC357 矩阵第K小【中等 堆 Java、Go、PHP】

题目

题目链接：
https://www.nowcoder.com/practice/c754e7a920614cba9b8b692ba9b20b5d

核心

堆，或者叫优先级队列
1

参考答案Java

import java.util.*;


public class Solution {
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     *
     * @param matrix int整型ArrayList>
     * @param k int整型
     * @return int整型
     */
    public int KthinMatrix (ArrayList<ArrayList<Integer>> matrix, int k) {
        //优先级队列
        int n = matrix.size(); //行和列都是n
        boolean[][] visited = new boolean[n][n];
        PriorityQueue<int[]> pq = new PriorityQueue<>(new Comparator<int[]>() {
            public int compare(int[] a,
                               int[] b) { //int[]:0:横坐标x  1：纵坐标y  2：xy对应的位置值
                return matrix.get(a[0]).get(a[1]) - matrix.get(b[0]).get(b[1]);
            }
        });

        int[] root = {0, 0, matrix.get(0).get(0)};

        visited[0][0] = true;
        pq.add(root);

        for (int i = 0; i < k - 1 ; i++) {
            int[] cur = pq.poll();
            int x = cur[0];
            int y = cur[1];

            if (x < n && y + 1 < matrix.get(x).size() && !visited[x][y + 1]) {
                visited[x][y + 1] = true;
                pq.add(new int[] {x, y + 1, matrix.get(x).get(y + 1)});
            }

            if (x + 1 < n && y < matrix.get(x).size() && !visited[x + 1][y]) {
                visited[x + 1][y] = true;
                pq.add(new int[] {x + 1, y, matrix.get(x + 1).get(y)});
            }
        }

        return pq.poll()[2];
    }

}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

参考答案Go

package main



/**
 * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
 *
 *
 * @param matrix int整型二维数组
 * @param k int整型
 * @return int整型
 */
func KthinMatrix(matrix [][]int, k int) int {
		//优先级队列，堆
	n := len(matrix)
	visited := make([][]bool, n)
	for i := 0; i < n; i++ {
		visited[i] = make([]bool, n)
	}

	pq := &MyPQ{make([]PQNode, 10), 0, 10}

	pq.add(0, 0, matrix[0][0])
	visited[0][0] = true

	for i := 0; i < k-1; i++ {
		poll := pq.poll()
		x := poll.x
		y := poll.y

		if x < n && y+1 < n && !visited[x][y+1] { //题目说了矩阵n*n
			pq.add(x, y+1, matrix[x][y+1])
			visited[x][y+1] = true
		}

		if x+1 < n && y < n && !visited[x+1][y] {
			pq.add(x+1, y, matrix[x+1][y])
			visited[x+1][y] = true
		}
	}

	return pq.poll().data
}

// 本答案用小顶堆，从上往下递增
type PQNode struct {
	x    int //横坐标
	y    int //纵坐标
	data int //坐标值
}
type MyPQ struct {
	arr         []PQNode
	size        int
	defaultSize int
}

func (pq *MyPQ) ensureCap(cap int) { //堆的扩容
	curlen := len(pq.arr)
	if curlen >= cap {
		return
	}

	newlen := curlen + curlen>>1
	newarr := make([]PQNode, newlen)
	for i := 0; i < curlen; i++ {
		newarr[i] = pq.arr[i]
	}

	pq.arr = newarr
}

func (pq *MyPQ) add(x, y, data int) {
	pq.ensureCap(pq.size + 1)
	node := PQNode{x, y, data}
	pq.arr[pq.size] = node
	pq.shiftup(pq.size)
	pq.size++

}

// 上滤
func (pq *MyPQ) shiftup(idx int) {
	root := pq.arr[idx]
	for idx > 0 {
		pidx := (idx - 1) >> 1
		pdata := pq.arr[pidx]

		if pdata.data < root.data {
			break
		}

		pq.arr[idx] = pdata
		idx = pidx
	}
	pq.arr[idx] = root
}

// 获取并删除堆顶元素，然后下窜，整理堆
func (pq *MyPQ) poll() PQNode {
	root := pq.arr[0]
	//fmt.Println("root:", root.data)
	pq.arr[0] = pq.arr[pq.size-1]
	pq.size -= 1
	pq.shiftdown(0)
	return root
}

func (pq *MyPQ) shiftdown(idx int) {
	root := pq.arr[idx]
	half := pq.size >> 1 //核心1: 用位移，别用除法

	for idx < half {
		childidx := (idx << 1) + 1 //核心2: 用位移，别用乘法
		rightidx := childidx + 1

		child := pq.arr[childidx]

		//小顶堆，哪个子节点更小，哪边就上去
		if rightidx < pq.size && pq.arr[childidx].data > pq.arr[rightidx].data {
			childidx = rightidx
			child = pq.arr[rightidx]
		}

		if child.data >= root.data {
			break
		}
		pq.arr[idx] = child
		idx = childidx
	}
	pq.arr[idx] = root
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

参考答案PHP

/**
 * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
 *
 * 
 * @param matrix int整型二维数组 
 * @param k int整型 
 * @return int整型
 */
function KthinMatrix( $matrix ,  $k )
{
    // 堆，优先级队列
    $pq = new PQ();
    $n =count($matrix);
    $visited = [];
    for($i=0;$i<$n;$i++){
        for($j=0;$j<$n;$j++){
            $visited[$i][$j] =false;
        }
    }

    $pq->add(0,0,$matrix[0][0]);
    $visited[0][0]= true;

    for($i=0;$i<$k-1;$i++){
        $poll = $pq->poll();
        $x = $poll[0];
        $y = $poll[1];

        if($x <$n && $y+1<$n && !$visited[$x][$y+1]){
            $pq->add($x,$y+1,$matrix[$x][$y+1]);
            $visited[$x][$y+1]=true;
        }

        if($x+1<$n&&$y<$n && !$visited[$x+1][$y]){
            $pq->add($x+1,$y,$matrix[$x+1][$y]);
            $visited[$x+1][$y]=true;
        }
    }

    return $pq->poll()[2];
}


//本答案用对，小顶对，也就是从上往下是递增的
class PQ{
    public $arr=[]; //二维数组，0：横坐标 1：纵坐标  2：坐标值
    public $defaultsize =10;
    public $size;
    public function __construct()
    {
        $this->size =0;
        for($i=0;$i<$this->defaultsize;$i++){
            $this->arr[$i] =[];
        }
    }


    public function ensurecap($cap){
        $curlen =count($this->arr);
        if($curlen >= $cap){
            return;
        }

        $newlen = $curlen+$curlen>>1;
        $newarr = [];
        for($i=0;$i<$newlen;$i++){
            if($i<$curlen){
                $newarr[$i] = $this->arr[$i];
            }else{
                $newarr[$i] =[];
            }
        }
        $this->arr = $newarr;
    }

    public function  add($x,$y,$data){
        $this->ensurecap($this->size+1);
        $this->arr[$this->size] = [$x,$y,$data];
        $this->shiftup($this->size);
        $this->size++;
    }


    //上滤
    public function shiftup($idx){
        $root = $this->arr[$idx];
        while ($idx>0){
            $pidx = ($idx-1) >>1;
            $parent= $this->arr[$pidx];

            if($parent[2] <=$root[2]){ //当前节点比父节点大，退出
                break;
            }

            $this->arr[$idx] = $parent;
            $idx = $pidx;
        }
        $this->arr[$idx] = $root;
    }

    public function poll() {
        $root = $this->arr[0];
        $this->arr[0] = $this->arr[$this->size-1];
        $this->size-=1;
        $this->shiftdown(0);
        return $root;
    }

    //下窜
    public function shiftdown($idx){
        $root = $this->arr[$idx];
        $half = $this->size >>1;
        while ($idx<$half){
            $childidx = ($idx << 1)+1;
            $rightidx = $childidx+1;
            $child = $this->arr[$childidx];

            if($rightidx<$this->size && $child[2] > $this->arr[$rightidx][2]){
                $childidx = $rightidx;
                $child = $this->arr[$rightidx];
            }

            if($child[2] >=$root[2]){
                break;
            }

            $this->arr[$idx] = $child;
            $idx =$childidx;
        }
        $this->arr[$idx] = $root;
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136